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 5, 02:15 UTC
[Quantization] Consolidate experts_int8 with fp8 online quantization (#38463)
Signed-off-by: Josephasafg <ajgard7@gmail.com>
Asaf Gardin committed
Apr 16, 2026 at 23:12 UTC
ad2b1277f99f1cee3d59d687aa1813267299f6a2
9 files changed
+403
-313
tests/quantization/test_experts_int8.py
-1
@@ -38,6 +38,5 @@ def test_model_experts_int8_startup(
38
dtype=dtype,
39
enforce_eager=True,
40
quantization="experts_int8",
41
- allow_deprecated_quantization=True,
41
) as vllm_model:
42
vllm_model.generate_greedy(example_prompts, max_tokens)
vllm/config/quantization.py
+4
@@ -19,6 +19,10 @@ class OnlineQuantScheme(Enum):
19
# blocks of 128x128 elements (popularized by DeepSeek)
20
FP8_PER_BLOCK = "fp8_per_block"
21
22
+ # int8, weight-only per-channel quantization for MoE expert weights.
23
+ # Linear layers remain unquantized.
24
+ INT8_PER_CHANNEL_WEIGHT_ONLY = "int8_per_channel_weight_only"
25
+
26
# TODO(future PRs): add more online quant schemes here: mxfp8, etc
27
28
vllm/model_executor/layers/fused_moe/oracle/int8.py
new
+84
@@ -0,0 +1,84 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+
6
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
7
+from vllm.logger import init_logger
8
+from vllm.model_executor.layers.fused_moe.all2all_utils import (
9
+ maybe_make_prepare_finalize,
10
+)
11
+from vllm.model_executor.layers.fused_moe.config import (
12
+ FusedMoEConfig,
13
+ FusedMoEQuantConfig,
14
+ int8_w8a16_moe_quant_config,
15
+)
16
+from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
17
+ SharedExperts,
18
+)
19
+
20
+logger = init_logger(__name__)
21
+
22
+
23
+def select_int8_moe_backend(
24
+ 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,
34
+ )
35
+ if not supported:
36
+ raise ValueError(
37
+ f"INT8 Triton MoE backend does not support the "
38
+ f"deployment configuration: {reason}"
39
+ )
40
+
41
+ logger.info_once("Using Triton INT8 MoE backend", scope="local")
42
+ return TritonExperts
43
+
44
+
45
+def make_int8_moe_quant_config(
46
+ w1_scale: torch.Tensor,
47
+ w2_scale: torch.Tensor,
48
+) -> FusedMoEQuantConfig:
49
+ return int8_w8a16_moe_quant_config(
50
+ w1_scale=w1_scale,
51
+ w2_scale=w2_scale,
52
+ w1_zp=None,
53
+ w2_zp=None,
54
+ )
55
+
56
+
57
+def make_int8_moe_kernel(
58
+ moe_quant_config: FusedMoEQuantConfig,
59
+ moe_config: FusedMoEConfig,
60
+ experts_cls: type[mk.FusedMoEExperts],
61
+ routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
62
+ shared_experts: SharedExperts | None = None,
63
+) -> mk.FusedMoEKernel:
64
+ prepare_finalize = maybe_make_prepare_finalize(
65
+ moe=moe_config,
66
+ quant_config=moe_quant_config,
67
+ routing_tables=routing_tables,
68
+ allow_new_interface=True,
69
+ )
70
+ assert prepare_finalize is not None
71
+
72
+ logger.info_once("Using %s", prepare_finalize.__class__.__name__, scope="local")
73
+
74
+ experts = experts_cls(
75
+ moe_config=moe_config,
76
+ quant_config=moe_quant_config,
77
+ )
78
+
79
+ return mk.FusedMoEKernel(
80
+ prepare_finalize,
81
+ experts,
82
+ shared_experts=shared_experts,
83
+ inplace=not moe_config.disable_inplace,
84
+ )
vllm/model_executor/layers/quantization/__init__.py
+1
-1
@@ -40,6 +40,7 @@ QuantizationMethods = Literal[
40
# shorthand for creating a more complicated online quant config object
41
"fp8_per_tensor",
42
"fp8_per_block",
43
+ "int8_per_channel_weight_only",
44
]
45
QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods))
46
@@ -47,7 +48,6 @@ DEPRECATED_QUANTIZATION_METHODS = [
48
"tpu_int8",
49
"fbgemm_fp8",
50
"fp_quant",
50
- "experts_int8",
51
]
52
53
# The customized quantization methods which will be added to this dict.
vllm/model_executor/layers/quantization/experts_int8.py
+11
-157
@@ -5,27 +5,25 @@ from typing import Any
5
6
import torch
7
8
-from vllm.distributed import get_tensor_model_parallel_rank, get_tp_group
9
-from vllm.model_executor.layers.fused_moe import (
10
- FusedMoE,
11
- FusedMoEConfig,
12
- FusedMoEMethodBase,
13
-)
14
-from vllm.model_executor.layers.fused_moe.config import (
15
- FusedMoEQuantConfig,
16
- int8_w8a16_moe_quant_config,
17
-)
8
+from vllm.model_executor.layers.fused_moe import FusedMoE
9
from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
10
from vllm.model_executor.layers.quantization import QuantizationMethods
11
from vllm.model_executor.layers.quantization.base_config import (
12
QuantizationConfig,
13
QuantizeMethodBase,
14
)
24
-from vllm.model_executor.utils import set_weight_attrs
15
+from vllm.model_executor.layers.quantization.online.int8 import (
16
+ Int8OnlineMoEMethod,
17
+)
18
19
20
class ExpertsInt8Config(QuantizationConfig):
28
- """Config class for Int8 experts quantization."""
21
+ """Online int8 quantization for MoE expert weights.
22
+ Linear layers are left unquantized.
23
+
24
+ Backward-compatible config for ``--quantization experts_int8``.
25
+ Prefer ``--quantization int8_per_channel``
26
+ """
27
28
def __init__(self) -> None:
29
super().__init__()
@@ -56,149 +54,5 @@ class ExpertsInt8Config(QuantizationConfig):
54
if isinstance(layer, LinearBase):
55
return UnquantizedLinearMethod()
56
elif isinstance(layer, FusedMoE):
59
- return ExpertsInt8MoEMethod(self, layer.moe_config)
57
+ return Int8OnlineMoEMethod(layer=layer)
58
return None
61
-
62
-
63
-class ExpertsInt8MoEMethod(FusedMoEMethodBase):
64
- def __init__(
65
- self,
66
- quant_config: ExpertsInt8Config,
67
- moe: FusedMoEConfig,
68
- ):
69
- super().__init__(moe)
70
- self.quant_config = quant_config
71
-
72
- def create_weights(
73
- self,
74
- layer: torch.nn.Module,
75
- num_experts: int,
76
- hidden_size: int,
77
- intermediate_size_per_partition: int,
78
- params_dtype: torch.dtype,
79
- **extra_weight_attrs,
80
- ):
81
- int8_dtype = torch.int8
82
-
83
- assert "weight_loader" in extra_weight_attrs
84
- weight_loader = extra_weight_attrs["weight_loader"]
85
- wrapped_weight_loader = ExpertsInt8MoEMethod.quantizing_weight_loader(
86
- layer, weight_loader
87
- )
88
- extra_weight_attrs["weight_loader"] = wrapped_weight_loader
89
-
90
- # Fused gate_up_proj (column parallel)
91
- w13_weight = torch.nn.Parameter(
92
- torch.empty(
93
- num_experts,
94
- 2 * intermediate_size_per_partition,
95
- hidden_size,
96
- dtype=int8_dtype,
97
- ),
98
- requires_grad=False,
99
- )
100
- layer.register_parameter("w13_weight", w13_weight)
101
- set_weight_attrs(w13_weight, extra_weight_attrs)
102
-
103
- # down_proj (row parallel)
104
- w2_weight = torch.nn.Parameter(
105
- torch.empty(
106
- num_experts,
107
- hidden_size,
108
- intermediate_size_per_partition,
109
- dtype=int8_dtype,
110
- ),
111
- requires_grad=False,
112
- )
113
- layer.register_parameter("w2_weight", w2_weight)
114
- set_weight_attrs(w2_weight, extra_weight_attrs)
115
-
116
- w13_scale = torch.nn.Parameter(
117
- torch.zeros(
118
- num_experts, 2 * intermediate_size_per_partition, dtype=torch.float32
119
- ),
120
- requires_grad=False,
121
- )
122
- layer.register_parameter("w13_scale", w13_scale)
123
-
124
- w2_scale = torch.nn.Parameter(
125
- torch.zeros(num_experts, hidden_size, dtype=torch.float32),
126
- requires_grad=False,
127
- )
128
- layer.register_parameter("w2_scale", w2_scale)
129
-
130
- def get_fused_moe_quant_config(
131
- self, layer: torch.nn.Module
132
- ) -> FusedMoEQuantConfig | None:
133
- return int8_w8a16_moe_quant_config(
134
- w1_scale=layer.w13_scale, w2_scale=layer.w2_scale, w1_zp=None, w2_zp=None
135
- )
136
-
137
- def apply(
138
- self,
139
- layer: FusedMoE,
140
- x: torch.Tensor,
141
- topk_weights: torch.Tensor,
142
- topk_ids: torch.Tensor,
143
- shared_experts_input: torch.Tensor | None,
144
- ) -> torch.Tensor:
145
- from vllm.model_executor.layers.fused_moe import fused_experts
146
-
147
- return fused_experts(
148
- x,
149
- layer.w13_weight,
150
- layer.w2_weight,
151
- topk_weights=topk_weights,
152
- topk_ids=topk_ids,
153
- inplace=not self.moe.disable_inplace,
154
- activation=layer.activation,
155
- apply_router_weight_on_input=layer.apply_router_weight_on_input,
156
- global_num_experts=layer.global_num_experts,
157
- expert_map=layer.expert_map,
158
- quant_config=self.moe_quant_config,
159
- )
160
-
161
- @staticmethod
162
- def quantizing_weight_loader(layer, weight_loader):
163
- def quantize_and_call_weight_loader(
164
- param: torch.nn.Parameter,
165
- loaded_weight: torch.Tensor,
166
- weight_name: str,
167
- shard_id: int,
168
- expert_id: int,
169
- ):
170
- tp_rank = get_tensor_model_parallel_rank()
171
- shard_size = layer.intermediate_size_per_partition
172
- shard = slice(tp_rank * shard_size, (tp_rank + 1) * shard_size)
173
- device = get_tp_group().device
174
- loaded_weight = loaded_weight.to(device)
175
- # w1, gate_proj case: Load into first shard of w13.
176
- if shard_id == "w1":
177
- scales = quantize_in_place_and_get_scales(loaded_weight[shard, :])
178
- layer.w13_scale.data[expert_id, 0:shard_size].copy_(scales[:, 0])
179
- # w3, up_proj case: Load into second shard of w13.
180
- elif shard_id == "w3":
181
- scales = quantize_in_place_and_get_scales(loaded_weight[shard, :])
182
- layer.w13_scale.data[expert_id, shard_size : 2 * shard_size].copy_(
183
- scales[:, 0]
184
- )
185
- # w2, down_proj case: Load into only shard of w2.
186
- elif shard_id == "w2":
187
- scales = quantize_in_place_and_get_scales(loaded_weight[:, shard])
188
- layer.w2_scale.data[expert_id, :].copy_(scales[:, 0])
189
- else:
190
- raise ValueError(f"Shard id must be in [0,1,2] but got {shard_id}")
191
- weight_loader(param, loaded_weight, weight_name, shard_id, expert_id)
192
-
193
- return quantize_and_call_weight_loader
194
-
195
-
196
-def quantize_in_place_and_get_scales(weight: torch.Tensor) -> torch.Tensor:
197
- vmax = torch.iinfo(torch.int8).max
198
- scales = torch.max(torch.abs(weight), dim=1, keepdim=True)[0] / vmax
199
-
200
- weight.div_(scales)
201
- weight.round_()
202
- weight.clamp_(-vmax, vmax)
203
-
204
- return scales
vllm/model_executor/layers/quantization/online/base.py
+16
-2
@@ -9,6 +9,7 @@ from vllm.config.quantization import (
9
OnlineQuantizationConfigArgs,
10
OnlineQuantScheme,
11
)
12
+from vllm.logger import init_logger
13
from vllm.model_executor.layers.fused_moe import (
14
FusedMoE,
15
)
@@ -33,6 +34,11 @@ from vllm.model_executor.layers.quantization.online.fp8 import (
34
Fp8PerTensorOnlineLinearMethod,
35
Fp8PerTensorOnlineMoEMethod,
36
)
37
+from vllm.model_executor.layers.quantization.online.int8 import (
38
+ Int8OnlineMoEMethod,
39
+)
40
+
41
+logger = init_logger(__name__)
42
43
44
class OnlineQuantizationConfig(QuantizationConfig):
@@ -96,7 +102,13 @@ class OnlineQuantizationConfig(QuantizationConfig):
102
return UnquantizedLinearMethod()
103
104
linear_scheme = self.args.linear_scheme_override or self.args.global_scheme
99
- if linear_scheme == OnlineQuantScheme.FP8_PER_BLOCK:
105
+ if linear_scheme == OnlineQuantScheme.INT8_PER_CHANNEL_WEIGHT_ONLY:
106
+ logger.warning_once(
107
+ "INT8 online quantization only quantizes MoE expert "
108
+ "weights. linear layers remain in full precision."
109
+ )
110
+ return UnquantizedLinearMethod()
111
+ elif linear_scheme == OnlineQuantScheme.FP8_PER_BLOCK:
112
return Fp8PerBlockOnlineLinearMethod()
113
else:
114
return Fp8PerTensorOnlineLinearMethod()
@@ -109,7 +121,9 @@ class OnlineQuantizationConfig(QuantizationConfig):
121
return UnquantizedFusedMoEMethod(layer.moe_config)
122
123
moe_scheme = self.args.moe_scheme_override or self.args.global_scheme
112
- if moe_scheme == OnlineQuantScheme.FP8_PER_BLOCK:
124
+ if moe_scheme == OnlineQuantScheme.INT8_PER_CHANNEL_WEIGHT_ONLY:
125
+ return Int8OnlineMoEMethod(layer=layer)
126
+ elif moe_scheme == OnlineQuantScheme.FP8_PER_BLOCK:
127
return Fp8PerBlockOnlineMoEMethod(layer=layer)
128
else:
129
return Fp8PerTensorOnlineMoEMethod(layer=layer)
vllm/model_executor/layers/quantization/online/fp8.py
+6
-152
@@ -10,7 +10,6 @@ if TYPE_CHECKING:
10
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11
from vllm.model_executor.layers.fused_moe import FusedMoE
12
from vllm.model_executor.layers.fused_moe.config import (
13
- FusedMoEConfig,
13
FusedMoEQuantConfig,
14
)
15
from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend
@@ -19,15 +18,15 @@ import vllm.envs as envs
18
from vllm import _custom_ops as ops
19
from vllm.config import get_current_vllm_config
20
from vllm.model_executor.kernels.linear import init_fp8_linear_kernel
22
-from vllm.model_executor.layers.fused_moe import (
23
- FusedMoEMethodBase,
24
-)
21
from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
22
select_fp8_moe_backend,
23
)
24
from vllm.model_executor.layers.linear import (
25
LinearMethodBase,
26
)
27
+from vllm.model_executor.layers.quantization.online.moe_base import (
28
+ OnlineMoEMethodBase,
29
+)
30
from vllm.model_executor.layers.quantization.utils.quant_utils import (
31
GroupShape,
32
create_fp8_quant_key,
@@ -44,7 +43,7 @@ from vllm.model_executor.model_loader.reload.layerwise import (
43
initialize_online_processing,
44
)
45
from vllm.model_executor.parameter import ModelWeightParameter
47
-from vllm.model_executor.utils import replace_parameter, set_weight_attrs
46
+from vllm.model_executor.utils import replace_parameter
47
from vllm.platforms import current_platform
48
from vllm.utils.deep_gemm import per_block_cast_to_fp8
49
@@ -268,21 +267,15 @@ class Fp8PerBlockOnlineLinearMethod(_Fp8OnlineLinearBase):
267
# ---------------------------------------------------------------------------
268
269
271
-class _Fp8OnlineMoEBase(FusedMoEMethodBase):
270
+class _Fp8OnlineMoEBase(OnlineMoEMethodBase):
271
"""Shared base for online FP8 MoE methods. Loads fp16/bf16 checkpoint
272
weights onto meta device and materializes them just-in-time."""
273
275
- uses_meta_device: bool = True
276
-
274
# Declared here for mypy; actual values are set in __init__.
275
fp8_backend: "Fp8MoeBackend"
276
experts_cls: "type[mk.FusedMoEExperts] | None"
277
weight_scale_name: str
278
weight_block_size: list[int] | None
282
- moe: "FusedMoEConfig"
283
- is_monolithic: bool
284
- moe_quant_config: "FusedMoEQuantConfig | None"
285
- moe_kernel: "mk.FusedMoEKernel | None"
279
280
def __init__(
281
self,
@@ -313,77 +306,6 @@ class _Fp8OnlineMoEBase(FusedMoEMethodBase):
306
allow_vllm_cutlass=False,
307
)
308
316
- def create_weights(
317
- self,
318
- layer: Module,
319
- num_experts: int,
320
- hidden_size: int,
321
- intermediate_size_per_partition: int,
322
- params_dtype: torch.dtype,
323
- **extra_weight_attrs,
324
- ):
325
- layer.num_experts = num_experts
326
- layer.orig_dtype = params_dtype
327
- layer.weight_block_size = None
328
-
329
- # WEIGHTS
330
- w13_weight = torch.nn.Parameter(
331
- torch.empty(
332
- num_experts,
333
- 2 * intermediate_size_per_partition,
334
- hidden_size,
335
- device="meta",
336
- dtype=params_dtype,
337
- ),
338
- requires_grad=False,
339
- )
340
- layer.register_parameter("w13_weight", w13_weight)
341
- set_weight_attrs(w13_weight, extra_weight_attrs)
342
-
343
- w2_weight = torch.nn.Parameter(
344
- torch.empty(
345
- num_experts,
346
- hidden_size,
347
- intermediate_size_per_partition,
348
- device="meta", # materialized and processed during loading
349
- dtype=params_dtype,
350
- ),
351
- requires_grad=False,
352
- )
353
- layer.register_parameter("w2_weight", w2_weight)
354
- set_weight_attrs(w2_weight, extra_weight_attrs)
355
-
356
- # BIASES (for models like GPT-OSS that have biased MoE)
357
- if self.moe.has_bias:
358
- w13_bias = torch.nn.Parameter(
359
- torch.zeros(
360
- num_experts,
361
- 2 * intermediate_size_per_partition,
362
- device="meta", # materialized and processed during loading
363
- dtype=layer.orig_dtype,
364
- ),
365
- requires_grad=False,
366
- )
367
- layer.register_parameter("w13_bias", w13_bias)
368
- set_weight_attrs(w13_bias, extra_weight_attrs)
369
-
370
- w2_bias = torch.nn.Parameter(
371
- torch.zeros(
372
- num_experts,
373
- hidden_size,
374
- device="meta", # materialized and processed during loading
375
- dtype=layer.orig_dtype,
376
- ),
377
- requires_grad=False,
378
- )
379
- layer.register_parameter("w2_bias", w2_bias)
380
- set_weight_attrs(w2_bias, extra_weight_attrs)
381
-
382
- layer.w13_input_scale = None
383
- layer.w2_input_scale = None
384
-
385
- initialize_online_processing(layer)
386
-
309
def _setup_kernel(
310
self,
311
layer: "FusedMoE",
@@ -430,15 +352,6 @@ class _Fp8OnlineMoEBase(FusedMoEMethodBase):
352
shared_experts=layer.shared_experts,
353
)
354
433
- def maybe_make_prepare_finalize(
434
- self,
435
- routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
436
- ) -> "mk.FusedMoEPrepareAndFinalizeModular | None":
437
- raise ValueError(
438
- f"{self.__class__.__name__} uses the new modular kernel "
439
- "initialization logic. This function should not be called."
440
- )
441
-
355
def get_fused_moe_quant_config(
356
self, layer: torch.nn.Module
357
) -> "FusedMoEQuantConfig":
@@ -460,68 +373,9 @@ class _Fp8OnlineMoEBase(FusedMoEMethodBase):
373
block_shape=self.weight_block_size,
374
)
375
463
- # Inject biases into the quant config if the model has them
464
- # (e.g. GPT-OSS biased MoE)
465
- if quant_config is not None and self.moe.has_bias:
466
- w13_bias = getattr(layer, "w13_bias", None)
467
- w2_bias = getattr(layer, "w2_bias", None)
468
- if w13_bias is not None:
469
- quant_config._w1.bias = w13_bias
470
- if w2_bias is not None:
471
- quant_config._w2.bias = w2_bias
472
-
376
+ self._maybe_inject_biases(quant_config, layer)
377
return quant_config
378
475
- @property
476
- def supports_eplb(self) -> bool:
477
- return True
478
-
479
- def apply_monolithic(
480
- self,
481
- layer: "FusedMoE",
482
- x: torch.Tensor,
483
- router_logits: torch.Tensor,
484
- ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
485
- assert self.is_monolithic
486
- assert self.moe_kernel is not None
487
- return self.moe_kernel.apply_monolithic(
488
- x,
489
- layer.w13_weight,
490
- layer.w2_weight,
491
- router_logits,
492
- activation=layer.activation,
493
- global_num_experts=layer.global_num_experts,
494
- expert_map=layer.expert_map,
495
- apply_router_weight_on_input=layer.apply_router_weight_on_input,
496
- num_expert_group=layer.num_expert_group,
497
- topk_group=layer.topk_group,
498
- e_score_correction_bias=layer.e_score_correction_bias,
499
- routed_scaling_factor=layer.routed_scaling_factor,
500
- )
501
-
502
- def apply(
503
- self,
504
- layer: "FusedMoE",
505
- x: torch.Tensor,
506
- topk_weights: torch.Tensor,
507
- topk_ids: torch.Tensor,
508
- shared_experts_input: torch.Tensor | None,
509
- ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
510
- assert not self.is_monolithic
511
- assert self.moe_kernel is not None
512
- return self.moe_kernel.apply(
513
- x,
514
- layer.w13_weight,
515
- layer.w2_weight,
516
- topk_weights,
517
- topk_ids,
518
- activation=layer.activation,
519
- global_num_experts=layer.global_num_experts,
520
- expert_map=layer.expert_map,
521
- apply_router_weight_on_input=layer.apply_router_weight_on_input,
522
- shared_experts_input=shared_experts_input,
523
- )
524
-
379
380
class Fp8PerTensorOnlineMoEMethod(_Fp8OnlineMoEBase):
381
"""Online tensorwise FP8 MoE quantization.
vllm/model_executor/layers/quantization/online/int8.py
new
+109
@@ -0,0 +1,109 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from typing import TYPE_CHECKING
5
+
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
11
+ from vllm.model_executor.layers.fused_moe import FusedMoE
12
+ from vllm.model_executor.layers.fused_moe.config import (
13
+ FusedMoEQuantConfig,
14
+ )
15
+
16
+from vllm.model_executor.layers.fused_moe.oracle.int8 import (
17
+ make_int8_moe_kernel,
18
+ make_int8_moe_quant_config,
19
+ select_int8_moe_backend,
20
+)
21
+from vllm.model_executor.layers.quantization.online.moe_base import (
22
+ OnlineMoEMethodBase,
23
+)
24
+from vllm.model_executor.utils import replace_parameter
25
+
26
+
27
+class Int8OnlineMoEMethod(OnlineMoEMethodBase):
28
+ """Online per-channel INT8 MoE quantization.
29
+ Loads fp16/bf16 weights and quantizes them per-row to int8 during loading.
30
+ """
31
+
32
+ def __init__(
33
+ self,
34
+ *,
35
+ layer: torch.nn.Module,
36
+ ):
37
+ super().__init__(layer.moe_config)
38
+ self.experts_cls: type[mk.FusedMoEExperts] = select_int8_moe_backend(
39
+ config=self.moe,
40
+ )
41
+
42
+ def process_weights_after_loading(self, layer: Module) -> None:
43
+ if getattr(layer, "_already_called_process_weights_after_loading", False):
44
+ return
45
+
46
+ self._quantize_weights(layer)
47
+ self._setup_kernel(layer)
48
+
49
+ layer._already_called_process_weights_after_loading = True
50
+
51
+ def _quantize_weights(self, layer: Module) -> None:
52
+ vmax = torch.iinfo(torch.int8).max
53
+
54
+ w13 = torch.empty_like(layer.w13_weight, dtype=torch.int8)
55
+ w2 = torch.empty_like(layer.w2_weight, dtype=torch.int8)
56
+ w13_scale = torch.zeros(
57
+ layer.num_experts,
58
+ layer.w13_weight.shape[1],
59
+ device=w13.device,
60
+ dtype=torch.float32,
61
+ )
62
+ w2_scale = torch.zeros(
63
+ layer.num_experts,
64
+ layer.w2_weight.shape[1],
65
+ device=w2.device,
66
+ dtype=torch.float32,
67
+ )
68
+
69
+ for expert in range(layer.local_num_experts):
70
+ # w13: per-row quantization over hidden_size dim
71
+ w = layer.w13_weight[expert, :, :]
72
+ scales = w.abs().amax(dim=1) / vmax
73
+ q = w.div(scales.unsqueeze(1)).round().clamp(-vmax, vmax)
74
+ w13[expert, :, :] = q.to(torch.int8)
75
+ w13_scale[expert, :] = scales
76
+
77
+ # w2: per-row quantization over intermediate_size dim
78
+ w = layer.w2_weight[expert, :, :]
79
+ scales = w.abs().amax(dim=1) / vmax
80
+ q = w.div(scales.unsqueeze(1)).round().clamp(-vmax, vmax)
81
+ w2[expert, :, :] = q.to(torch.int8)
82
+ w2_scale[expert, :] = scales
83
+
84
+ replace_parameter(layer, "w13_weight", w13)
85
+ replace_parameter(layer, "w2_weight", w2)
86
+ replace_parameter(layer, "w13_scale", w13_scale)
87
+ replace_parameter(layer, "w2_scale", w2_scale)
88
+
89
+ def _setup_kernel(self, layer: "FusedMoE") -> None:
90
+ self.moe_quant_config = self.get_fused_moe_quant_config(layer)
91
+ assert self.moe_quant_config is not None
92
+ assert self.experts_cls is not None
93
+ self.moe_kernel = make_int8_moe_kernel(
94
+ moe_quant_config=self.moe_quant_config,
95
+ moe_config=self.moe,
96
+ experts_cls=self.experts_cls,
97
+ routing_tables=layer._maybe_init_expert_routing_tables(),
98
+ shared_experts=layer.shared_experts,
99
+ )
100
+
101
+ def get_fused_moe_quant_config(
102
+ self, layer: torch.nn.Module
103
+ ) -> "FusedMoEQuantConfig | None":
104
+ quant_config = make_int8_moe_quant_config(
105
+ w1_scale=layer.w13_scale,
106
+ w2_scale=layer.w2_scale,
107
+ )
108
+ self._maybe_inject_biases(quant_config, layer)
109
+ return quant_config
vllm/model_executor/layers/quantization/online/moe_base.py
new
+172
@@ -0,0 +1,172 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from abc import abstractmethod
5
+
6
+import torch
7
+
8
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9
+from vllm.model_executor.layers.fused_moe import FusedMoEMethodBase
10
+from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
11
+from vllm.model_executor.model_loader.reload.layerwise import (
12
+ initialize_online_processing,
13
+)
14
+from vllm.model_executor.utils import set_weight_attrs
15
+
16
+
17
+class OnlineMoEMethodBase(FusedMoEMethodBase):
18
+ """Base for MoE methods that load full-precision weights on meta device
19
+ and quantize them after loading via the QeRL layerwise processing system.
20
+ """
21
+
22
+ uses_meta_device: bool = True
23
+
24
+ def create_weights(
25
+ self,
26
+ layer: torch.nn.Module,
27
+ num_experts: int,
28
+ hidden_size: int,
29
+ intermediate_size_per_partition: int,
30
+ params_dtype: torch.dtype,
31
+ **extra_weight_attrs,
32
+ ):
33
+ layer.num_experts = num_experts
34
+ layer.orig_dtype = params_dtype
35
+ layer.weight_block_size = None
36
+
37
+ # Fused gate_up_proj (column parallel) — full precision on meta device
38
+ w13_weight = torch.nn.Parameter(
39
+ torch.empty(
40
+ num_experts,
41
+ 2 * intermediate_size_per_partition,
42
+ hidden_size,
43
+ device="meta",
44
+ dtype=params_dtype,
45
+ ),
46
+ requires_grad=False,
47
+ )
48
+ layer.register_parameter("w13_weight", w13_weight)
49
+ set_weight_attrs(w13_weight, extra_weight_attrs)
50
+
51
+ # down_proj (row parallel) — full precision on meta device
52
+ w2_weight = torch.nn.Parameter(
53
+ torch.empty(
54
+ num_experts,
55
+ hidden_size,
56
+ intermediate_size_per_partition,
57
+ device="meta",
58
+ dtype=params_dtype,
59
+ ),
60
+ requires_grad=False,
61
+ )
62
+ layer.register_parameter("w2_weight", w2_weight)
63
+ set_weight_attrs(w2_weight, extra_weight_attrs)
64
+
65
+ # BIASES (for models like GPT-OSS that have biased MoE)
66
+ if self.moe.has_bias:
67
+ w13_bias = torch.nn.Parameter(
68
+ torch.zeros(
69
+ num_experts,
70
+ 2 * intermediate_size_per_partition,
71
+ device="meta",
72
+ dtype=layer.orig_dtype,
73
+ ),
74
+ requires_grad=False,
75
+ )
76
+ layer.register_parameter("w13_bias", w13_bias)
77
+ set_weight_attrs(w13_bias, extra_weight_attrs)
78
+
79
+ w2_bias = torch.nn.Parameter(
80
+ torch.zeros(
81
+ num_experts,
82
+ hidden_size,
83
+ device="meta",
84
+ dtype=layer.orig_dtype,
85
+ ),
86
+ requires_grad=False,
87
+ )
88
+ layer.register_parameter("w2_bias", w2_bias)
89
+ set_weight_attrs(w2_bias, extra_weight_attrs)
90
+
91
+ layer.w13_input_scale = None
92
+ layer.w2_input_scale = None
93
+
94
+ initialize_online_processing(layer)
95
+
96
+ @abstractmethod
97
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
98
+ pass
99
+
100
+ def _maybe_inject_biases(
101
+ self,
102
+ quant_config: FusedMoEQuantConfig,
103
+ layer: torch.nn.Module,
104
+ ) -> None:
105
+ """Inject biases into the quant config if the model has them
106
+ (e.g. GPT-OSS biased MoE)."""
107
+ if self.moe.has_bias:
108
+ w13_bias = getattr(layer, "w13_bias", None)
109
+ w2_bias = getattr(layer, "w2_bias", None)
110
+ if w13_bias is not None:
111
+ quant_config._w1.bias = w13_bias
112
+ if w2_bias is not None:
113
+ quant_config._w2.bias = w2_bias
114
+
115
+ def maybe_make_prepare_finalize(
116
+ self,
117
+ routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
118
+ ) -> mk.FusedMoEPrepareAndFinalizeModular | None:
119
+ raise ValueError(
120
+ f"{self.__class__.__name__} uses the new modular kernel "
121
+ "initialization logic. This function should not be called."
122
+ )
123
+
124
+ @property
125
+ def supports_eplb(self) -> bool:
126
+ return True
127
+
128
+ def apply_monolithic(
129
+ self,
130
+ layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
131
+ x: torch.Tensor,
132
+ router_logits: torch.Tensor,
133
+ ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
134
+ assert self.is_monolithic
135
+ assert self.moe_kernel is not None
136
+ return self.moe_kernel.apply_monolithic(
137
+ x,
138
+ layer.w13_weight,
139
+ layer.w2_weight,
140
+ router_logits,
141
+ activation=layer.activation,
142
+ global_num_experts=layer.global_num_experts,
143
+ expert_map=layer.expert_map,
144
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
145
+ num_expert_group=layer.num_expert_group,
146
+ topk_group=layer.topk_group,
147
+ e_score_correction_bias=layer.e_score_correction_bias,
148
+ routed_scaling_factor=layer.routed_scaling_factor,
149
+ )
150
+
151
+ def apply(
152
+ self,
153
+ layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
154
+ x: torch.Tensor,
155
+ topk_weights: torch.Tensor,
156
+ topk_ids: torch.Tensor,
157
+ shared_experts_input: torch.Tensor | None,
158
+ ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
159
+ assert not self.is_monolithic
160
+ assert self.moe_kernel is not None
161
+ return self.moe_kernel.apply(
162
+ x,
163
+ layer.w13_weight,
164
+ layer.w2_weight,
165
+ topk_weights,
166
+ topk_ids,
167
+ activation=layer.activation,
168
+ global_num_experts=layer.global_num_experts,
169
+ expert_map=layer.expert_map,
170
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
171
+ shared_experts_input=shared_experts_input,
172
+ )