@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 5f76b3fb30
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 4, 20:00 UTC

[MoE] Convert CT W8A8 To Oracle Structure (#39187)

Signed-off-by: Robert Shaw <robertgshaw2@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>

Robert Shaw committed Apr 22, 2026 at 10:53 UTC 5f76b3fb3044785e628384f41fd3b32f1185b448
6 files changed +248 -63
vllm/model_executor/layers/fused_moe/fused_batched_moe.py
+10 -10
@@ -928,16 +928,16 @@ class BatchedTritonExperts(mk.FusedMoEExpertsModular):
928 p.is_cuda() and p.has_device_capability((8, 9))
929 )
930
931 - SUPPORTED_W_A_FP8 = [
932 - (kFp8Static128BlockSym, kFp8Dynamic128Sym),
933 - (kFp8StaticChannelSym, kFp8DynamicTokenSym),
934 - (kFp8StaticTensorSym, kFp8DynamicTokenSym),
935 - (kFp8StaticTensorSym, kFp8StaticTensorSym),
936 - (kFp8StaticTensorSym, kFp8DynamicTensorSym),
937 - ]
938 - return (weight_key, activation_key) == (None, None) or (
939 - device_supports_fp8 and (weight_key, activation_key) in SUPPORTED_W_A_FP8
940 - )
931 + supported: list[tuple[QuantKey | None, QuantKey | None]] = [(None, None)]
932 + if device_supports_fp8:
933 + supported += [
934 + (kFp8Static128BlockSym, kFp8Dynamic128Sym),
935 + (kFp8StaticChannelSym, kFp8DynamicTokenSym),
936 + (kFp8StaticTensorSym, kFp8DynamicTokenSym),
937 + (kFp8StaticTensorSym, kFp8StaticTensorSym),
938 + (kFp8StaticTensorSym, kFp8DynamicTensorSym),
939 + ]
940 + return (weight_key, activation_key) in supported
941
942 @staticmethod
943 def _supports_activation(activation: MoEActivation) -> bool:
vllm/model_executor/layers/fused_moe/fused_moe.py
+20 -12
@@ -46,6 +46,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
46 kFp8Static128BlockSym,
47 kFp8StaticChannelSym,
48 kFp8StaticTensorSym,
49 + kInt8DynamicTokenSym,
50 + kInt8StaticChannelSym,
51 )
52 from vllm.platforms import current_platform
53 from vllm.triton_utils import tl, triton
@@ -1952,18 +1954,24 @@ class TritonExperts(mk.FusedMoEExpertsModular):
1954 weight_key: QuantKey | None,
1955 activation_key: QuantKey | None,
1956 ) -> bool:
1955 - if not current_platform.supports_fp8():
1956 - return (weight_key, activation_key) == (None, None)
1957 -
1958 - SUPPORTED_W_A = [
1959 - (None, None),
1960 - (kFp8Static128BlockSym, kFp8Dynamic128Sym),
1961 - (kFp8StaticChannelSym, kFp8DynamicTokenSym),
1962 - (kFp8StaticTensorSym, kFp8DynamicTokenSym),
1963 - (kFp8StaticTensorSym, kFp8StaticTensorSym),
1964 - (kFp8StaticTensorSym, kFp8DynamicTensorSym),
1965 - ]
1966 - return (weight_key, activation_key) in SUPPORTED_W_A
1957 + # INT8 requires at least 7.5 (Turing).
1958 + device_supports_int8 = (
1959 + current_platform.is_cuda()
1960 + and current_platform.has_device_capability((7, 5))
1961 + )
1962 +
1963 + supported: list[tuple[QuantKey | None, QuantKey | None]] = [(None, None)]
1964 + if device_supports_int8:
1965 + supported.append((kInt8StaticChannelSym, kInt8DynamicTokenSym))
1966 + if current_platform.supports_fp8():
1967 + supported += [
1968 + (kFp8Static128BlockSym, kFp8Dynamic128Sym),
1969 + (kFp8StaticChannelSym, kFp8DynamicTokenSym),
1970 + (kFp8StaticTensorSym, kFp8DynamicTokenSym),
1971 + (kFp8StaticTensorSym, kFp8StaticTensorSym),
1972 + (kFp8StaticTensorSym, kFp8DynamicTensorSym),
1973 + ]
1974 + return (weight_key, activation_key) in supported
1975
1976 @staticmethod
1977 def _supports_activation(activation: MoEActivation) -> bool:
vllm/model_executor/layers/fused_moe/oracle/int8.py
+160 -23
@@ -1,9 +1,12 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +from enum import Enum
5 +
6 import torch
7
8 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9 +from vllm.config.kernel import MoEBackend
10 from vllm.logger import init_logger
11 from vllm.model_executor.layers.fused_moe.all2all_utils import (
12 maybe_make_prepare_finalize,
@@ -11,46 +14,165 @@ from vllm.model_executor.layers.fused_moe.all2all_utils import (
14 from vllm.model_executor.layers.fused_moe.config import (
15 FusedMoEConfig,
16 FusedMoEQuantConfig,
17 + int8_w8a8_moe_quant_config,
18 int8_w8a16_moe_quant_config,
19 )
20 from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
21 SharedExperts,
22 )
23 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
24 + QuantKey,
25 + kInt8DynamicTokenSym,
26 + kInt8StaticChannelSym,
27 +)
28
29 logger = init_logger(__name__)
30
31
32 +class Int8MoeBackend(Enum):
33 + TRITON = "TRITON"
34 +
35 +
36 +def _get_priority_backends(
37 + moe_config: FusedMoEConfig,
38 +) -> list[Int8MoeBackend]:
39 + """
40 + Get available backends in priority order based on platform and config.
41 + """
42 + return [Int8MoeBackend.TRITON]
43 +
44 +
45 +def backend_to_kernel_cls(
46 + backend: Int8MoeBackend,
47 +) -> list[type[mk.FusedMoEExperts]]:
48 + if backend == Int8MoeBackend.TRITON:
49 + from vllm.model_executor.layers.fused_moe.fused_moe import (
50 + TritonExperts,
51 + )
52 +
53 + return [TritonExperts]
54 +
55 + else:
56 + raise ValueError(f"Unknown Int8 MoE backend: {backend.value}")
57 +
58 +
59 +def map_int8_backend(runner_backend: MoEBackend) -> Int8MoeBackend:
60 + """Map user's MoEBackend to Int8MoeBackend."""
61 + mapping = {
62 + "triton": Int8MoeBackend.TRITON,
63 + }
64 + if backend := mapping.get(runner_backend):
65 + return backend
66 + raise ValueError(
67 + f"moe_backend='{runner_backend}' is not supported for Int8 MoE. "
68 + f"Expected one of {list(mapping.keys())}."
69 + )
70 +
71 +
72 def select_int8_moe_backend(
73 config: FusedMoEConfig,
25 -) -> type[mk.FusedMoEExperts]:
26 - from vllm.model_executor.layers.fused_moe.fused_moe import TritonExperts
27 -
28 - supported, reason = TritonExperts.is_supported_config(
29 - TritonExperts,
30 - config,
31 - None,
32 - None,
33 - mk.FusedMoEActivationFormat.Standard,
74 + weight_key: QuantKey | None = kInt8StaticChannelSym,
75 + activation_key: QuantKey | None = kInt8DynamicTokenSym,
76 +) -> tuple[Int8MoeBackend, type[mk.FusedMoEExperts]]:
77 + """
78 + Select the primary Int8 MoE backend.
79 + Note: Shape-specific fallbacks may still occur at runtime.
80 + """
81 +
82 + if config.is_lora_enabled:
83 + return Int8MoeBackend.TRITON, backend_to_kernel_cls(Int8MoeBackend.TRITON)[0]
84 +
85 + AVAILABLE_BACKENDS = _get_priority_backends(config)
86 +
87 + activation_format = (
88 + mk.FusedMoEActivationFormat.BatchedExperts
89 + if config.moe_parallel_config.use_batched_activation_format
90 + else mk.FusedMoEActivationFormat.Standard
91 )
35 - if not supported:
36 - raise ValueError(
37 - f"INT8 Triton MoE backend does not support the "
38 - f"deployment configuration: {reason}"
92 +
93 + def _make_log_backend(backend: Int8MoeBackend) -> str:
94 + available_backend_strs = [b.value for b in AVAILABLE_BACKENDS]
95 + return (
96 + f"Using {backend.value} Int8 MoE backend out "
97 + f"of potential backends: {available_backend_strs}."
98 )
99
41 - logger.info_once("Using Triton INT8 MoE backend", scope="local")
42 - return TritonExperts
100 + def _make_log_unsupported(backend: Int8MoeBackend, reason: str | None) -> str:
101 + if reason:
102 + return (
103 + f"Int8 MoE backend {backend.value} does not support the "
104 + f"deployment configuration since {reason}."
105 + )
106 + else:
107 + return (
108 + f"Int8 MoE backend '{backend.value}' does not support the "
109 + "deployment configuration."
110 + )
111 +
112 + def _return_or_raise(
113 + backend: Int8MoeBackend,
114 + ) -> tuple[Int8MoeBackend, type[mk.FusedMoEExperts]]:
115 + for k_cls in backend_to_kernel_cls(backend):
116 + supported, reason = k_cls.is_supported_config(
117 + k_cls, config, weight_key, activation_key, activation_format
118 + )
119 + if supported:
120 + logger.info_once(_make_log_backend(backend), scope="local")
121 + return backend, k_cls
122 + raise ValueError(_make_log_unsupported(backend, reason))
123 +
124 + # Handle explicit moe_backend from user.
125 + runner_backend = config.moe_backend
126 + if runner_backend != "auto":
127 + requested_backend = map_int8_backend(runner_backend)
128 + return _return_or_raise(requested_backend)
129 +
130 + # Select kernels in order of backend.
131 + for backend in AVAILABLE_BACKENDS:
132 + for k_cls in backend_to_kernel_cls(backend):
133 + supported, reason = k_cls.is_supported_config(
134 + k_cls,
135 + config,
136 + weight_key,
137 + activation_key,
138 + activation_format,
139 + )
140 + if supported:
141 + logger.info_once(_make_log_backend(backend), scope="local")
142 + return backend, k_cls
143 + else:
144 + logger.debug_once(_make_log_unsupported(backend, reason), scope="local")
145 +
146 + raise NotImplementedError(
147 + "No Int8 MoE backend supports the deployment configuration."
148 + )
149
150
151 def make_int8_moe_quant_config(
152 w1_scale: torch.Tensor,
153 w2_scale: torch.Tensor,
154 + a1_scale: torch.Tensor | None = None,
155 + a2_scale: torch.Tensor | None = None,
156 + per_act_token_quant: bool = False,
157 ) -> FusedMoEQuantConfig:
49 - return int8_w8a16_moe_quant_config(
158 + assert (a1_scale is None and a2_scale is None) or (
159 + a1_scale is not None and a2_scale is not None
160 + ), "a1_scale and a2_scale must both be provided or both be None"
161 +
162 + if a1_scale is None or a2_scale is None:
163 + return int8_w8a16_moe_quant_config(
164 + w1_scale=w1_scale,
165 + w2_scale=w2_scale,
166 + w1_zp=None,
167 + w2_zp=None,
168 + )
169 +
170 + return int8_w8a8_moe_quant_config(
171 w1_scale=w1_scale,
172 w2_scale=w2_scale,
52 - w1_zp=None,
53 - w2_zp=None,
173 + a1_scale=a1_scale,
174 + a2_scale=a2_scale,
175 + per_act_token_quant=per_act_token_quant,
176 )
177
178
@@ -61,24 +183,39 @@ def make_int8_moe_kernel(
183 routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
184 shared_experts: SharedExperts | None = None,
185 ) -> mk.FusedMoEKernel:
186 + # Create Prepare/Finalize.
187 prepare_finalize = maybe_make_prepare_finalize(
188 moe=moe_config,
189 quant_config=moe_quant_config,
190 routing_tables=routing_tables,
191 allow_new_interface=True,
192 + use_monolithic=issubclass(experts_cls, mk.FusedMoEExpertsMonolithic),
193 )
194 assert prepare_finalize is not None
195
196 logger.info_once("Using %s", prepare_finalize.__class__.__name__, scope="local")
197
74 - experts = experts_cls(
75 - moe_config=moe_config,
76 - quant_config=moe_quant_config,
77 - )
198 + # Create Experts.
199 + if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts:
200 + max_num_tokens = prepare_finalize.max_num_tokens_per_rank()
201 + assert max_num_tokens is not None
202 + experts = experts_cls(
203 + moe_config=moe_config,
204 + quant_config=moe_quant_config,
205 + max_num_tokens=max_num_tokens,
206 + num_dispatchers=prepare_finalize.num_dispatchers(),
207 + )
208 + else:
209 + experts = experts_cls(
210 + moe_config=moe_config,
211 + quant_config=moe_quant_config,
212 + )
213
79 - return mk.FusedMoEKernel(
214 + kernel = mk.FusedMoEKernel(
215 prepare_finalize,
216 experts,
217 shared_experts=shared_experts,
218 inplace=not moe_config.disable_inplace,
219 )
220 +
221 + return kernel
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_int8.py
+48 -16
@@ -8,6 +8,7 @@ from compressed_tensors.quantization import (
8 QuantizationStrategy,
9 )
10
11 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
12 from vllm.logger import init_logger
13 from vllm.model_executor.layers.fused_moe import (
14 FusedMoE,
@@ -16,17 +17,27 @@ from vllm.model_executor.layers.fused_moe import (
17 from vllm.model_executor.layers.fused_moe.config import (
18 FusedMoEConfig,
19 FusedMoEQuantConfig,
19 - int8_w8a8_moe_quant_config,
20 +)
21 +from vllm.model_executor.layers.fused_moe.oracle.int8 import (
22 + make_int8_moe_kernel,
23 + make_int8_moe_quant_config,
24 + select_int8_moe_backend,
25 )
26 from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501
27 CompressedTensorsMoEMethod,
28 )
29 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
30 + kInt8DynamicTokenSym,
31 + kInt8StaticChannelSym,
32 +)
33 from vllm.model_executor.utils import set_weight_attrs
34
35 logger = init_logger(__name__)
36
37
38 class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
39 + """W8A8 Int8 MoE quantization using compressed tensors."""
40 +
41 def __init__(
42 self,
43 weight_quant: QuantizationArgs,
@@ -56,6 +67,13 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
67 "dynamic per token quantization. Found static input scales."
68 )
69
70 + # Select Int8 MoE backend.
71 + self.int8_backend, self.experts_cls = select_int8_moe_backend(
72 + config=self.moe,
73 + weight_key=kInt8StaticChannelSym,
74 + activation_key=kInt8DynamicTokenSym,
75 + )
76 +
77 def create_weights(
78 self,
79 layer: torch.nn.Module,
@@ -122,13 +140,28 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
140 layer.w13_input_scale = None
141 layer.w2_input_scale = None
142
125 - def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
126 - pass
143 + def process_weights_after_loading(self, layer: FusedMoE) -> None:
144 + self.moe_quant_config = self.get_fused_moe_quant_config(layer)
145 + assert self.experts_cls is not None
146 + self.moe_kernel = make_int8_moe_kernel(
147 + moe_quant_config=self.moe_quant_config,
148 + moe_config=self.moe,
149 + experts_cls=self.experts_cls,
150 + routing_tables=layer._maybe_init_expert_routing_tables(),
151 + shared_experts=layer.shared_experts,
152 + )
153 +
154 + def maybe_make_prepare_finalize(
155 + self,
156 + routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
157 + ) -> mk.FusedMoEPrepareAndFinalizeModular | None:
158 + raise ValueError(
159 + f"{self.__class__.__name__} uses the new modular kernel initialization "
160 + "logic. This function should not be called."
161 + )
162
128 - def get_fused_moe_quant_config(
129 - self, layer: torch.nn.Module
130 - ) -> FusedMoEQuantConfig | None:
131 - return int8_w8a8_moe_quant_config(
163 + def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
164 + return make_int8_moe_quant_config(
165 w1_scale=layer.w13_weight_scale,
166 w2_scale=layer.w2_weight_scale,
167 a1_scale=layer.w13_input_scale,
@@ -144,18 +177,17 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
177 topk_ids: torch.Tensor,
178 shared_experts_input: torch.Tensor | None,
179 ) -> torch.Tensor:
147 - from vllm.model_executor.layers.fused_moe import fused_experts
148 -
149 - return fused_experts(
150 - hidden_states=x,
151 - w1=layer.w13_weight,
152 - w2=layer.w2_weight,
180 + assert not self.is_monolithic
181 + assert self.moe_kernel is not None
182 + return self.moe_kernel.apply(
183 + x,
184 + layer.w13_weight,
185 + layer.w2_weight,
186 topk_weights=topk_weights,
187 topk_ids=topk_ids,
155 - inplace=not self.moe.disable_inplace,
188 activation=layer.activation,
157 - apply_router_weight_on_input=layer.apply_router_weight_on_input,
189 global_num_experts=layer.global_num_experts,
190 expert_map=layer.expert_map,
160 - quant_config=self.moe_quant_config,
191 + apply_router_weight_on_input=layer.apply_router_weight_on_input,
192 + shared_experts_input=shared_experts_input,
193 )
vllm/model_executor/layers/quantization/online/int8.py
+7 -2
@@ -7,7 +7,6 @@ import torch
7 from torch.nn import Module
8
9 if TYPE_CHECKING:
10 - import vllm.model_executor.layers.fused_moe.modular_kernel as mk
10 from vllm.model_executor.layers.fused_moe import FusedMoE
11 from vllm.model_executor.layers.fused_moe.config import (
12 FusedMoEQuantConfig,
@@ -21,6 +20,10 @@ from vllm.model_executor.layers.fused_moe.oracle.int8 import (
20 from vllm.model_executor.layers.quantization.online.moe_base import (
21 OnlineMoEMethodBase,
22 )
23 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
24 + kInt8DynamicTokenSym,
25 + kInt8StaticChannelSym,
26 +)
27 from vllm.model_executor.utils import replace_parameter
28
29
@@ -35,8 +38,10 @@ class Int8OnlineMoEMethod(OnlineMoEMethodBase):
38 layer: torch.nn.Module,
39 ):
40 super().__init__(layer.moe_config)
38 - self.experts_cls: type[mk.FusedMoEExperts] = select_int8_moe_backend(
41 + self.int8_backend, self.experts_cls = select_int8_moe_backend(
42 config=self.moe,
43 + weight_key=kInt8StaticChannelSym,
44 + activation_key=kInt8DynamicTokenSym,
45 )
46
47 def process_weights_after_loading(self, layer: Module) -> None:
vllm/model_executor/layers/quantization/utils/quant_utils.py
+3
@@ -170,6 +170,9 @@ kMxfp8Dynamic = QuantKey(FP8_DTYPE, scale=kMxfp8DynamicGroupScale, symmetric=Tru
170 kMxfp4StaticGroupScale = ScaleDesc(MXFP_SCALE_DTYPE, True, GroupShape(1, 32))
171 kMxfp4Static = QuantKey(FP4_DTYPE, scale=kMxfp4StaticGroupScale, symmetric=True)
172
173 +kInt8StaticChannelSym = QuantKey(torch.int8, kStaticChannelScale, symmetric=True)
174 +kInt8DynamicTokenSym = QuantKey(torch.int8, kDynamicTokenScale, symmetric=True)
175 +
176
177 def create_fp8_quant_key(
178 static: bool,