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 2, 23:15 UTC
[MoE] Migrate W4A8 CT to oracle kernel setup (#42680)
Signed-off-by: Siddharth Bedekar <bedeksid@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Siddharth Bedekar committed
May 23, 2026 at 10:56 UTC
4438b6e7dc480dd59e5edabfcc939c15321a129a
4 files changed
+328
-330
vllm/model_executor/layers/fused_moe/__init__.py
-2
@@ -93,7 +93,6 @@ if HAS_TRITON:
93
CutlassBatchedExpertsFp8,
94
CutlassExpertsFp8,
95
CutlassExpertsW4A8Fp8,
96
- cutlass_moe_w4a8_fp8,
96
)
97
from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import (
98
DeepGemmExperts,
@@ -133,7 +132,6 @@ if HAS_TRITON:
132
"fused_experts",
133
"get_config_file_name",
134
"GroupedTopk",
136
- "cutlass_moe_w4a8_fp8",
135
"CutlassExpertsFp8",
136
"CutlassBatchedExpertsFp8",
137
"CutlassExpertsW4A8Fp8",
vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py
+54
-147
@@ -20,9 +20,6 @@ from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
20
moe_permute,
21
moe_unpermute,
22
)
23
-from vllm.model_executor.layers.fused_moe.prepare_finalize import (
24
- MoEPrepareAndFinalizeNoDPEPModular,
25
-)
23
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
24
TopKWeightAndReduceDelegate,
25
TopKWeightAndReduceNoOP,
@@ -36,6 +33,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
33
kFp8DynamicTokenSym,
34
kFp8StaticChannelSym,
35
kFp8StaticTensorSym,
36
+ kInt4Static,
37
kMxfp4Dynamic,
38
kMxfp4Static,
39
kNvfp4Dynamic,
@@ -1237,29 +1235,36 @@ def run_cutlass_moe_w4a8_fp8(
1235
class CutlassExpertsW4A8Fp8(mk.FusedMoEExpertsModular):
1236
def __init__(
1237
self,
1240
- out_dtype: torch.dtype | None,
1241
- a_strides1: torch.Tensor,
1242
- a_strides2: torch.Tensor,
1243
- b_strides1: torch.Tensor,
1244
- b_strides2: torch.Tensor,
1245
- c_strides1: torch.Tensor,
1246
- c_strides2: torch.Tensor,
1247
- s_strides1: torch.Tensor,
1248
- s_strides2: torch.Tensor,
1238
moe_config: FusedMoEConfig,
1239
quant_config: FusedMoEQuantConfig,
1240
+ b_strides1: torch.Tensor,
1241
+ b_strides2: torch.Tensor,
1242
group_size: int,
1243
):
1244
super().__init__(moe_config=moe_config, quant_config=quant_config)
1254
- self.out_dtype = out_dtype
1255
- self.a_strides1 = a_strides1
1256
- self.a_strides2 = a_strides2
1245
+
1246
+ e = moe_config.num_local_experts
1247
+ n = moe_config.intermediate_size_per_partition
1248
+ k = moe_config.hidden_dim
1249
+ device = moe_config.device
1250
+
1251
+ self.out_dtype = moe_config.in_dtype
1252
+
1253
+ a_strides1_c_strides2 = torch.full((e,), k, device=device, dtype=torch.int64)
1254
+ self.a_strides1 = a_strides1_c_strides2
1255
+ self.a_strides2 = torch.full((e,), n, device=device, dtype=torch.int64)
1256
+ self.c_strides1 = torch.full((e,), 2 * n, device=device, dtype=torch.int64)
1257
+ self.c_strides2 = a_strides1_c_strides2
1258
+
1259
self.b_strides1 = b_strides1
1260
self.b_strides2 = b_strides2
1259
- self.c_strides1 = c_strides1
1260
- self.c_strides2 = c_strides2
1261
- self.s_strides1 = s_strides1
1262
- self.s_strides2 = s_strides2
1261
+
1262
+ # sizeof(StrideS) = 16 bytes, encoded as 2xint64.
1263
+ self.s_strides1 = torch.zeros((e, 2), device=device, dtype=torch.int64)
1264
+ self.s_strides1[:, 0] = 2 * n
1265
+ self.s_strides2 = torch.zeros((e, 2), device=device, dtype=torch.int64)
1266
+ self.s_strides2[:, 0] = k
1267
+
1268
self.group_size = group_size
1269
1270
@staticmethod
@@ -1267,42 +1272,53 @@ class CutlassExpertsW4A8Fp8(mk.FusedMoEExpertsModular):
1272
return mk.FusedMoEActivationFormat.Standard
1273
1274
@staticmethod
1270
- def _supports_current_device() -> bool:
1271
- raise NotImplementedError(
1272
- "CutlassExpertsW4A8Fp8 is not yet used by an Oracle. "
1273
- "This method should not be called."
1275
+ def is_supported_config(
1276
+ cls: type[mk.FusedMoEExperts],
1277
+ moe_config: FusedMoEConfig,
1278
+ weight_key: QuantKey | None,
1279
+ activation_key: QuantKey | None,
1280
+ activation_format: mk.FusedMoEActivationFormat,
1281
+ ) -> tuple[bool, str | None]:
1282
+ if moe_config.in_dtype != torch.bfloat16:
1283
+ return (
1284
+ False,
1285
+ f"kernel does not support {moe_config.in_dtype} input/output dtype",
1286
+ )
1287
+
1288
+ return mk.FusedMoEExperts.is_supported_config(
1289
+ cls,
1290
+ moe_config,
1291
+ weight_key,
1292
+ activation_key,
1293
+ activation_format,
1294
)
1295
1296
+ @staticmethod
1297
+ def _supports_current_device() -> bool:
1298
+ return cutlass_group_gemm_supported()
1299
+
1300
@staticmethod
1301
def _supports_no_act_and_mul() -> bool:
1278
- raise NotImplementedError(
1279
- "CutlassExpertsW4A8Fp8 is not yet used by an Oracle. "
1280
- "This method should not be called."
1281
- )
1302
+ return False
1303
1304
@staticmethod
1305
def _supports_quant_scheme(
1306
weight_key: QuantKey | None,
1307
activation_key: QuantKey | None,
1308
) -> bool:
1288
- raise NotImplementedError(
1289
- "CutlassExpertsW4A8Fp8 is not yet used by an Oracle. "
1290
- "This method should not be called."
1291
- )
1309
+ return (weight_key, activation_key) == (kInt4Static, kFp8DynamicTokenSym)
1310
1311
@staticmethod
1312
def _supports_activation(activation: MoEActivation) -> bool:
1295
- raise NotImplementedError(
1296
- "CutlassExpertsW4A8Fp8 is not yet used by an Oracle. "
1297
- "This method should not be called."
1313
+ return activation in (
1314
+ MoEActivation.SILU,
1315
+ MoEActivation.GELU,
1316
+ MoEActivation.SWIGLUOAI,
1317
)
1318
1319
@staticmethod
1320
def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
1302
- raise NotImplementedError(
1303
- "CutlassExpertsW4A8Fp8 is not yet used by an Oracle. "
1304
- "This method should not be called."
1305
- )
1321
+ return True
1322
1323
def supports_expert_map(self) -> bool:
1324
return True
@@ -1394,112 +1410,3 @@ class CutlassExpertsW4A8Fp8(mk.FusedMoEExpertsModular):
1410
topk_weights,
1411
self.group_size,
1412
)
1397
-
1398
-
1399
-def cutlass_moe_w4a8_fp8(
1400
- a: torch.Tensor,
1401
- w1_q: torch.Tensor,
1402
- w2_q: torch.Tensor,
1403
- topk_weights: torch.Tensor,
1404
- topk_ids: torch.Tensor,
1405
- a_strides1: torch.Tensor,
1406
- a_strides2: torch.Tensor,
1407
- b_strides1: torch.Tensor,
1408
- b_strides2: torch.Tensor,
1409
- c_strides1: torch.Tensor,
1410
- c_strides2: torch.Tensor,
1411
- s_strides1: torch.Tensor,
1412
- s_strides2: torch.Tensor,
1413
- quant_config: FusedMoEQuantConfig,
1414
- moe_config: FusedMoEConfig,
1415
- activation: MoEActivation = MoEActivation.SILU,
1416
- expert_map: torch.Tensor | None = None,
1417
- apply_router_weight_on_input: bool = False,
1418
- global_num_experts: int = -1,
1419
- group_size: int = 128,
1420
-) -> torch.Tensor:
1421
- """
1422
- This function computes a w4a8-quantized Mixture of Experts (MoE) layer
1423
- using two sets of quantized weights, w1_q and w2_q, and top-k gating
1424
- mechanism. The matrix multiplications are implemented with CUTLASS
1425
- mixed-dtype grouped gemm.
1426
-
1427
- Parameters:
1428
- - a (torch.Tensor): The input tensor to the MoE layer.
1429
- Shape: [M, K]
1430
- - w1_q (torch.Tensor): The first set of fp8-quantized expert weights.
1431
- Shape: [num_experts, 2*N, K // packed_factor]
1432
- - w2_q (torch.Tensor): The second set of fp8-quantized expert weights.
1433
- Shape: [num_experts, K, N // packed_factor]
1434
- - topk_weights (torch.Tensor): The weights of each token->expert mapping.
1435
- - topk_ids (torch.Tensor): The token->expert mappings.
1436
- - a_strides1 (torch.Tensor): The input strides for the first gemm.
1437
- Shape: [num_experts]
1438
- - a_strides2 (torch.Tensor): The input strides for the second gemm.
1439
- Shape: [num_experts]
1440
- - b_strides1 (torch.Tensor): The packed layout for the first gemm weights.
1441
- Shape: [num_experts, 3]
1442
- dtype: torch.int32
1443
- - b_strides2 (torch.Tensor): The packed layout for the second gemm weights.
1444
- Shape: [num_experts, 3]
1445
- dtype: torch.int32
1446
- - c_strides1 (torch.Tensor): The output strides for the first gemm.
1447
- Shape: [num_experts]
1448
- - c_strides2 (torch.Tensor): The output strides for the second gemm.
1449
- Shape: [num_experts]
1450
- - s_strides1 (torch.Tensor): strides for the group-wise scales for the first gemm.
1451
- Shape: [num_experts, 2]
1452
- dtype: torch.int64
1453
- - s_strides2 (torch.Tensor): strides for the group-wise scales for the second gemm.
1454
- Shape: [num_experts, 2]
1455
- dtype: torch.int64
1456
- - per_act_token (Optional[bool]): Whether the scale is per-token or
1457
- per-tensor.
1458
- - activation (MoEActivation): The activation function to use.
1459
- - expert_map (Optional[torch.Tensor]): In the case of Expert parallel,
1460
- every Rank is responsible for a subset of experts. expert_map is a
1461
- mapping from global expert-id to local expert-id. When expert_map[i]
1462
- is -1, it means that this Rank is not responsible for global
1463
- expert-id i.
1464
- - apply_router_weight_on_input (bool): When true, the topk weights are
1465
- applied directly on the inputs. This is only applicable when topk is 1.
1466
- - global_num_experts (int): The total number of experts.
1467
- - group_size (int): The number of weights per scale factor
1468
-
1469
- Returns:
1470
- - torch.Tensor: The bf16 output tensor after applying the MoE layer.
1471
- """
1472
- assert quant_config is not None
1473
-
1474
- num_experts = global_num_experts if global_num_experts != -1 else w1_q.size(0)
1475
-
1476
- fn = mk.FusedMoEKernel(
1477
- MoEPrepareAndFinalizeNoDPEPModular(),
1478
- CutlassExpertsW4A8Fp8(
1479
- out_dtype=a.dtype,
1480
- a_strides1=a_strides1,
1481
- a_strides2=a_strides2,
1482
- b_strides1=b_strides1,
1483
- b_strides2=b_strides2,
1484
- c_strides1=c_strides1,
1485
- c_strides2=c_strides2,
1486
- s_strides1=s_strides1,
1487
- s_strides2=s_strides2,
1488
- moe_config=moe_config,
1489
- quant_config=quant_config,
1490
- group_size=group_size,
1491
- ),
1492
- inplace=False,
1493
- )
1494
-
1495
- return fn.apply(
1496
- a,
1497
- w1_q,
1498
- w2_q,
1499
- topk_weights,
1500
- topk_ids,
1501
- activation=activation,
1502
- global_num_experts=num_experts,
1503
- expert_map=expert_map,
1504
- apply_router_weight_on_input=apply_router_weight_on_input,
1505
- )
vllm/model_executor/layers/fused_moe/oracle/w4a8.py
new
+196
@@ -0,0 +1,196 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from enum import Enum
4
+from typing import TYPE_CHECKING
5
+
6
+import torch
7
+
8
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9
+from vllm.logger import init_logger
10
+from vllm.model_executor.layers.fused_moe.all2all_utils import (
11
+ maybe_make_prepare_finalize,
12
+)
13
+from vllm.model_executor.layers.fused_moe.config import (
14
+ FusedMoEConfig,
15
+ FusedMoEQuantConfig,
16
+ int4_w4afp8_moe_quant_config,
17
+)
18
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
19
+ QuantKey,
20
+ kFp8DynamicTokenSym,
21
+ kInt4Static,
22
+)
23
+
24
+if TYPE_CHECKING:
25
+ from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
26
+ CutlassExpertsW4A8Fp8,
27
+ )
28
+
29
+logger = init_logger(__name__)
30
+
31
+
32
+class W4A8MoeBackend(Enum):
33
+ CUTLASS = "CUTLASS"
34
+
35
+
36
+def backend_to_kernel_cls(
37
+ backend: W4A8MoeBackend,
38
+) -> list[type["CutlassExpertsW4A8Fp8"]]:
39
+ if backend == W4A8MoeBackend.CUTLASS:
40
+ from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
41
+ CutlassExpertsW4A8Fp8,
42
+ )
43
+
44
+ return [CutlassExpertsW4A8Fp8]
45
+ else:
46
+ raise ValueError(f"Unknown W4A8 MoE backend: {backend.value}")
47
+
48
+
49
+def select_w4a8_moe_backend(
50
+ config: FusedMoEConfig,
51
+ weight_key: QuantKey | None = kInt4Static,
52
+ activation_key: QuantKey | None = kFp8DynamicTokenSym,
53
+) -> tuple[W4A8MoeBackend, type["CutlassExpertsW4A8Fp8"]]:
54
+ backend = W4A8MoeBackend.CUTLASS
55
+
56
+ activation_format = (
57
+ mk.FusedMoEActivationFormat.BatchedExperts
58
+ if config.moe_parallel_config.use_batched_activation_format
59
+ else mk.FusedMoEActivationFormat.Standard
60
+ )
61
+
62
+ last_reason: str | None = None
63
+ for kernel_cls in backend_to_kernel_cls(backend):
64
+ supported, reason = kernel_cls.is_supported_config(
65
+ kernel_cls,
66
+ config,
67
+ weight_key,
68
+ activation_key,
69
+ activation_format,
70
+ )
71
+ if supported:
72
+ logger.info_once("Using %s W4A8 MoE backend.", backend.value)
73
+ return backend, kernel_cls
74
+ last_reason = reason
75
+
76
+ raise NotImplementedError(
77
+ f"W4A8 MoE backend {backend.value} does not support the "
78
+ f"deployment configuration: {last_reason}."
79
+ )
80
+
81
+
82
+def convert_to_w4a8_moe_kernel_format(
83
+ w13_weight_packed: torch.Tensor,
84
+ w2_weight_packed: torch.Tensor,
85
+ w13_weight_scale: torch.Tensor,
86
+ w2_weight_scale: torch.Tensor,
87
+) -> tuple[
88
+ torch.Tensor,
89
+ torch.Tensor,
90
+ torch.Tensor,
91
+ torch.Tensor,
92
+ torch.Tensor,
93
+ torch.Tensor,
94
+ torch.Tensor,
95
+ torch.Tensor,
96
+]:
97
+ from vllm import _custom_ops as ops
98
+ from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
99
+ from vllm.model_executor.layers.quantization.utils.quant_utils import (
100
+ GroupShape,
101
+ convert_bf16_scales_to_fp8,
102
+ convert_packed_uint4b8_to_signed_int4_inplace,
103
+ )
104
+
105
+ quant_fp8 = QuantFP8(static=False, group_shape=GroupShape.PER_TOKEN)
106
+
107
+ convert_packed_uint4b8_to_signed_int4_inplace(w13_weight_packed)
108
+ # Mirror the sync in CutlassW4A8LinearKernel; required for TP>1 correctness.
109
+ torch.accelerator.synchronize()
110
+ w13_weight_shuffled, b_strides1 = ops.cutlass_encode_and_reorder_int4b_grouped(
111
+ w13_weight_packed
112
+ )
113
+
114
+ convert_packed_uint4b8_to_signed_int4_inplace(w2_weight_packed)
115
+ # Mirror the sync in CutlassW4A8LinearKernel; required for TP>1 correctness.
116
+ torch.accelerator.synchronize()
117
+ w2_weight_shuffled, b_strides2 = ops.cutlass_encode_and_reorder_int4b_grouped(
118
+ w2_weight_packed
119
+ )
120
+
121
+ w13_weight_scale, w13_weight_chan_scale = convert_bf16_scales_to_fp8(
122
+ quant_fp8, w13_weight_scale
123
+ )
124
+ w2_weight_scale, w2_weight_chan_scale = convert_bf16_scales_to_fp8(
125
+ quant_fp8, w2_weight_scale
126
+ )
127
+
128
+ # Scales are stored as (E, N, K // 128), but the kernel expects
129
+ # (E, K // 128, N) in row-major format.
130
+ w13_weight_scale_packed = ops.cutlass_pack_scale_fp8(
131
+ w13_weight_scale.permute(0, 2, 1).contiguous()
132
+ )
133
+ w2_weight_scale_packed = ops.cutlass_pack_scale_fp8(
134
+ w2_weight_scale.permute(0, 2, 1).contiguous()
135
+ )
136
+
137
+ return (
138
+ w13_weight_shuffled,
139
+ w2_weight_shuffled,
140
+ w13_weight_scale_packed,
141
+ w2_weight_scale_packed,
142
+ w13_weight_chan_scale,
143
+ w2_weight_chan_scale,
144
+ b_strides1,
145
+ b_strides2,
146
+ )
147
+
148
+
149
+def make_w4a8_moe_quant_config(
150
+ w1_scale: torch.Tensor,
151
+ w2_scale: torch.Tensor,
152
+ g1_alphas: torch.Tensor,
153
+ g2_alphas: torch.Tensor,
154
+) -> FusedMoEQuantConfig:
155
+ return int4_w4afp8_moe_quant_config(
156
+ w1_scale=w1_scale,
157
+ w2_scale=w2_scale,
158
+ g1_alphas=g1_alphas,
159
+ g2_alphas=g2_alphas,
160
+ per_act_token_quant=True,
161
+ per_out_ch_quant=True,
162
+ )
163
+
164
+
165
+def make_w4a8_moe_kernel(
166
+ moe_quant_config: FusedMoEQuantConfig,
167
+ moe_config: FusedMoEConfig,
168
+ experts_cls: type["CutlassExpertsW4A8Fp8"],
169
+ b_strides1: torch.Tensor,
170
+ b_strides2: torch.Tensor,
171
+ group_size: int,
172
+ routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
173
+) -> mk.FusedMoEKernel:
174
+ prepare_finalize = maybe_make_prepare_finalize(
175
+ moe=moe_config,
176
+ quant_config=moe_quant_config,
177
+ routing_tables=routing_tables,
178
+ allow_new_interface=True,
179
+ )
180
+ assert prepare_finalize is not None
181
+
182
+ logger.info_once("Using %s", prepare_finalize.__class__.__name__)
183
+
184
+ experts = experts_cls(
185
+ moe_config=moe_config,
186
+ quant_config=moe_quant_config,
187
+ b_strides1=b_strides1,
188
+ b_strides2=b_strides2,
189
+ group_size=group_size,
190
+ )
191
+
192
+ return mk.FusedMoEKernel(
193
+ prepare_finalize,
194
+ experts,
195
+ inplace=not moe_config.disable_inplace,
196
+ )
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_fp8.py
+78
-181
@@ -7,12 +7,7 @@ from compressed_tensors.quantization import (
7
QuantizationArgs,
8
)
9
10
-import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11
-from vllm import _custom_ops as ops
12
-from vllm.logger import init_logger
10
from vllm.model_executor.layers.fused_moe import (
14
- FusedMoEActivationFormat,
15
- FusedMoEExpertsModular,
11
FusedMoeWeightScaleSupported,
12
RoutedExperts,
13
SharedExperts,
@@ -20,19 +15,18 @@ from vllm.model_executor.layers.fused_moe import (
15
from vllm.model_executor.layers.fused_moe.config import (
16
FusedMoEConfig,
17
FusedMoEQuantConfig,
23
- int4_w4afp8_moe_quant_config,
18
+)
19
+from vllm.model_executor.layers.fused_moe.oracle.w4a8 import (
20
+ convert_to_w4a8_moe_kernel_format,
21
+ make_w4a8_moe_kernel,
22
+ make_w4a8_moe_quant_config,
23
+ select_w4a8_moe_backend,
24
)
25
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501
26
CompressedTensorsMoEMethod,
27
)
28
-from vllm.model_executor.layers.quantization.utils.quant_utils import (
29
- convert_bf16_scales_to_fp8,
30
- convert_packed_uint4b8_to_signed_int4_inplace,
31
-)
28
from vllm.model_executor.utils import replace_parameter, set_weight_attrs
29
34
-logger = init_logger(__name__)
35
-
30
31
class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):
32
def __init__(
@@ -56,16 +50,10 @@ class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):
50
assert self.weight_quant.actorder != "group"
51
assert self.group_size == 128, "Only group size 128 supported for W4A8 MoE"
52
59
- self.disable_expert_map = False
60
- self.layer_name = layer_name
61
-
62
- from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
63
- from vllm.model_executor.layers.quantization.utils.quant_utils import (
64
- GroupShape,
53
+ self.w4a8_backend, self.experts_cls = select_w4a8_moe_backend(
54
+ config=self.moe,
55
)
56
67
- self.quant_fp8 = QuantFP8(static=False, group_shape=GroupShape.PER_TOKEN)
68
-
57
def create_weights(
58
self,
59
layer: torch.nn.Module,
@@ -156,152 +144,74 @@ class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):
144
layer.register_parameter("w13_weight_shape", w13_weight_shape)
145
set_weight_attrs(w13_weight_shape, extra_weight_attrs)
146
159
- # don't use input scales
160
- layer.w13_input_scale = None
161
- layer.w2_input_scale = None
162
-
163
- def process_weights_after_loading(self, layer):
164
- device = layer.w13_weight_packed.device
165
-
166
- # STRIDES
167
- # A, C
168
- self.a_strides1_c_strides2 = torch.full(
169
- (layer.local_num_experts,),
170
- layer.hidden_size,
171
- device=device,
172
- dtype=torch.int64,
173
- )
174
- self.a_strides2 = torch.full(
175
- (layer.local_num_experts,),
176
- layer.intermediate_size_per_partition,
177
- device=device,
178
- dtype=torch.int64,
179
- )
180
- self.c_strides1 = torch.full(
181
- (layer.local_num_experts,),
182
- 2 * layer.intermediate_size_per_partition,
183
- device=device,
184
- dtype=torch.int64,
185
- )
186
-
187
- # S (group-wise scales)
188
- # sizeof(StrideS) = 16 bytes, so we need to use 2xint64 to encode it
189
- self.s_strides1 = torch.zeros(
190
- (layer.local_num_experts, 2), device=device, dtype=torch.int64
191
- )
192
- self.s_strides1[:, 0] = 2 * layer.intermediate_size_per_partition
193
-
194
- self.s_strides2 = torch.zeros(
195
- (layer.local_num_experts, 2), device=device, dtype=torch.int64
196
- )
197
- self.s_strides2[:, 0] = layer.hidden_size
198
-
199
- # encode and reorder weight tensors, and get the layout to pass to
200
- # the grouped gemm kernel. `b_strides1/2` specifies the entire layout
201
- convert_packed_uint4b8_to_signed_int4_inplace(layer.w13_weight_packed)
202
- # mirror the sync in CutlassW4A8LinearKernel; required for tp>1 correctness
203
- torch.accelerator.synchronize()
204
- w13_weight_shuffled, self.b_strides1 = (
205
- ops.cutlass_encode_and_reorder_int4b_grouped(layer.w13_weight_packed)
206
- )
207
- replace_parameter(layer, "w13_weight_packed", w13_weight_shuffled)
208
- convert_packed_uint4b8_to_signed_int4_inplace(layer.w2_weight_packed)
209
- # mirror the sync in CutlassW4A8LinearKernel; required for tp>1 correctness
210
- torch.accelerator.synchronize()
211
- w2_weight_shuffled, self.b_strides2 = (
212
- ops.cutlass_encode_and_reorder_int4b_grouped(layer.w2_weight_packed)
213
- )
214
- replace_parameter(layer, "w2_weight_packed", w2_weight_shuffled)
215
-
216
- # convert bf16 scales to (fp8_scales, channel_scales)
217
- w13_weight_scale, w13_weight_chan_scale = convert_bf16_scales_to_fp8(
218
- self.quant_fp8, layer.w13_weight_scale
219
- )
220
- w2_weight_scale, w2_weight_chan_scale = convert_bf16_scales_to_fp8(
221
- self.quant_fp8, layer.w2_weight_scale
222
- )
223
-
224
- # register channel scales
225
- layer.register_parameter(
226
- "w13_weight_chan_scale",
227
- torch.nn.Parameter(w13_weight_chan_scale, requires_grad=False),
228
- )
229
- layer.register_parameter(
230
- "w2_weight_chan_scale",
231
- torch.nn.Parameter(w2_weight_chan_scale, requires_grad=False),
147
+ w13_weight_chan_scale = torch.nn.Parameter(
148
+ torch.ones(
149
+ num_experts,
150
+ 2 * intermediate_size_per_partition,
151
+ dtype=torch.float32,
152
+ ),
153
+ requires_grad=False,
154
)
155
+ layer.register_parameter("w13_weight_chan_scale", w13_weight_chan_scale)
156
234
- # The scales are stored as (E, N, K // 128) but the kernel expects
235
- # (E, K // 128, N) in row-major format, so we need to permute the last 2 dims
236
- # and make it contiguous
237
- w13_weight_scale_packed = ops.cutlass_pack_scale_fp8(
238
- w13_weight_scale.permute(0, 2, 1).contiguous()
239
- )
240
- replace_parameter(layer, "w13_weight_scale", w13_weight_scale_packed)
241
- w2_weight_scale_packed = ops.cutlass_pack_scale_fp8(
242
- w2_weight_scale.permute(0, 2, 1).contiguous()
157
+ w2_weight_chan_scale = torch.nn.Parameter(
158
+ torch.ones(num_experts, hidden_size, dtype=torch.float32),
159
+ requires_grad=False,
160
)
244
- replace_parameter(layer, "w2_weight_scale", w2_weight_scale_packed)
161
+ layer.register_parameter("w2_weight_chan_scale", w2_weight_chan_scale)
162
246
- def maybe_make_prepare_finalize(
247
- self,
248
- routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
249
- ) -> mk.FusedMoEPrepareAndFinalizeModular | None:
250
- return super().maybe_make_prepare_finalize(routing_tables)
163
+ # don't use input scales
164
+ layer.w13_input_scale = None
165
+ layer.w2_input_scale = None
166
252
- def get_fused_moe_quant_config(
253
- self, layer: torch.nn.Module
254
- ) -> FusedMoEQuantConfig | None:
255
- # Store quantization scales; both per-group and per-channel
256
- # Note we haven't specified the group size here because
257
- # the quant config logic assumes group-wise scaling
258
- # and channel-wise scaling are exclusive.
259
- return int4_w4afp8_moe_quant_config(
260
- w1_scale=layer.w13_weight_scale, # group scale
261
- w2_scale=layer.w2_weight_scale, # group scale
167
+ def process_weights_after_loading(self, layer: RoutedExperts) -> None:
168
+ (
169
+ w13_weight_packed,
170
+ w2_weight_packed,
171
+ w13_weight_scale,
172
+ w2_weight_scale,
173
+ w13_weight_chan_scale,
174
+ w2_weight_chan_scale,
175
+ b_strides1,
176
+ b_strides2,
177
+ ) = convert_to_w4a8_moe_kernel_format(
178
+ w13_weight_packed=layer.w13_weight_packed,
179
+ w2_weight_packed=layer.w2_weight_packed,
180
+ w13_weight_scale=layer.w13_weight_scale,
181
+ w2_weight_scale=layer.w2_weight_scale,
182
+ )
183
+
184
+ replace_parameter(layer, "w13_weight_packed", w13_weight_packed)
185
+ replace_parameter(layer, "w2_weight_packed", w2_weight_packed)
186
+ replace_parameter(layer, "w13_weight_scale", w13_weight_scale)
187
+ replace_parameter(layer, "w2_weight_scale", w2_weight_scale)
188
+ replace_parameter(layer, "w13_weight_chan_scale", w13_weight_chan_scale)
189
+ replace_parameter(layer, "w2_weight_chan_scale", w2_weight_chan_scale)
190
+
191
+ self.b_strides1 = b_strides1
192
+ self.b_strides2 = b_strides2
193
+
194
+ self.moe_quant_config = self.get_fused_moe_quant_config(layer)
195
+ if self.moe_quant_config is not None:
196
+ assert self.experts_cls is not None
197
+ self.moe_kernel = make_w4a8_moe_kernel(
198
+ moe_quant_config=self.moe_quant_config,
199
+ moe_config=self.moe,
200
+ experts_cls=self.experts_cls,
201
+ b_strides1=self.b_strides1,
202
+ b_strides2=self.b_strides2,
203
+ group_size=self.group_size,
204
+ routing_tables=layer._expert_routing_tables(),
205
+ )
206
+
207
+ def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
208
+ return make_w4a8_moe_quant_config(
209
+ w1_scale=layer.w13_weight_scale,
210
+ w2_scale=layer.w2_weight_scale,
211
g1_alphas=layer.w13_weight_chan_scale,
212
g2_alphas=layer.w2_weight_chan_scale,
264
- per_act_token_quant=True, # always use dynamic per-token
265
- per_out_ch_quant=True, # always use per-channel
266
- )
267
-
268
- def select_gemm_impl(
269
- self,
270
- prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular,
271
- layer: torch.nn.Module,
272
- ) -> mk.FusedMoEExpertsModular:
273
- assert self.moe_quant_config is not None
274
- assert (
275
- prepare_finalize.activation_format == FusedMoEActivationFormat.Standard
276
- ), "BatchedExperts not supported"
277
-
278
- from vllm.model_executor.layers.fused_moe import CutlassExpertsW4A8Fp8
279
-
280
- experts: FusedMoEExpertsModular
281
-
282
- logger.debug("CutlassExpertsW4A8Fp8(%s)", self.__class__.__name__)
283
- experts = CutlassExpertsW4A8Fp8(
284
- out_dtype=self.moe.in_dtype,
285
- a_strides1=self.a_strides1_c_strides2,
286
- a_strides2=self.a_strides2,
287
- b_strides1=self.b_strides1,
288
- b_strides2=self.b_strides2,
289
- c_strides1=self.c_strides1,
290
- c_strides2=self.a_strides1_c_strides2,
291
- s_strides1=self.s_strides1,
292
- s_strides2=self.s_strides2,
293
- moe_config=self.moe,
294
- quant_config=self.moe_quant_config,
295
- group_size=self.group_size,
213
)
214
298
- num_dispatchers = prepare_finalize.num_dispatchers()
299
- self.disable_expert_map = (
300
- num_dispatchers > 1 or not experts.supports_expert_map()
301
- )
302
-
303
- return experts
304
-
215
def apply(
216
self,
217
layer: RoutedExperts,
@@ -311,33 +221,20 @@ class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):
221
shared_experts: SharedExperts | None,
222
shared_experts_input: torch.Tensor | None,
223
) -> torch.Tensor:
314
- assert self.moe_quant_config is not None
315
-
316
- from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
317
- cutlass_moe_w4a8_fp8,
318
- )
319
-
320
- return cutlass_moe_w4a8_fp8(
321
- x,
322
- layer.w13_weight_packed,
323
- layer.w2_weight_packed,
324
- topk_weights,
325
- topk_ids,
326
- moe_config=self.moe,
327
- quant_config=self.moe_quant_config,
224
+ assert not self.is_monolithic
225
+ assert self.moe_kernel is not None
226
+ return self.moe_kernel.apply(
227
+ hidden_states=x,
228
+ w1=layer.w13_weight_packed,
229
+ w2=layer.w2_weight_packed,
230
+ topk_weights=topk_weights,
231
+ topk_ids=topk_ids,
232
activation=layer.activation,
233
global_num_experts=layer.global_num_experts,
330
- expert_map=None if self.disable_expert_map else layer.expert_map,
331
- a_strides1=self.a_strides1_c_strides2,
332
- a_strides2=self.a_strides2,
333
- b_strides1=self.b_strides1,
334
- b_strides2=self.b_strides2,
335
- c_strides1=self.c_strides1,
336
- c_strides2=self.a_strides1_c_strides2,
337
- s_strides1=self.s_strides1,
338
- s_strides2=self.s_strides2,
339
- group_size=self.group_size,
234
+ expert_map=layer.expert_map,
235
apply_router_weight_on_input=layer.apply_router_weight_on_input,
236
+ shared_experts=shared_experts,
237
+ shared_experts_input=shared_experts_input,
238
)
239
240
@property