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 10, 21:30 UTC
[MoE Refactor][12/N] Marlin Fp8 MoE Pure Function (#31499)
Signed-off-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com>
Robert Shaw committed
Dec 29, 2025 at 16:27 UTC
9152a30d8f6f85cbc5505fafb6f6682d4e189884
4 files changed
+92
-76
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+18
-5
@@ -964,12 +964,25 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
964
layer.w2_weight = torch.nn.Parameter(shuffled_w2, requires_grad=False)
965
966
elif self.use_marlin:
967
- prepare_moe_fp8_layer_for_marlin(
968
- layer, False, input_dtype=self.marlin_input_dtype
967
+ (
968
+ workspace,
969
+ w13_weight,
970
+ w2_weight,
971
+ w13_weight_scale,
972
+ w2_weight_scale,
973
+ ) = prepare_moe_fp8_layer_for_marlin(
974
+ layer,
975
+ layer.w13_weight,
976
+ layer.w2_weight,
977
+ layer.w13_weight_scale,
978
+ layer.w2_weight_scale,
979
+ input_dtype=self.marlin_input_dtype,
980
)
970
- # Activations not quantized for marlin.
971
- del layer.w13_input_scale
972
- del layer.w2_input_scale
981
+ layer.workspace = workspace
982
+ replace_parameter(layer, "w13_weight", w13_weight)
983
+ replace_parameter(layer, "w2_weight", w2_weight)
984
+ replace_parameter(layer, "w13_weight_scale", w13_weight_scale)
985
+ replace_parameter(layer, "w2_weight_scale", w2_weight_scale)
986
987
if self.use_cutlass:
988
assert self.weight_quant.strategy != QuantizationStrategy.BLOCK
vllm/model_executor/layers/quantization/fp8.py
+21
-21
@@ -912,6 +912,23 @@ class Fp8MoEMethod(FusedMoEMethodBase):
912
w13_weight, w2_weight = rocm_aiter_ops.shuffle_weights(
913
w13_weight, w2_weight
914
)
915
+ elif self.fp8_backend == Fp8MoeBackend.MARLIN:
916
+ (
917
+ workspace,
918
+ w13_weight,
919
+ w2_weight,
920
+ w13_weight_scale,
921
+ w2_weight_scale,
922
+ ) = prepare_moe_fp8_layer_for_marlin(
923
+ layer,
924
+ w13_weight,
925
+ w2_weight,
926
+ w13_weight_scale,
927
+ w2_weight_scale,
928
+ input_dtype=self.marlin_input_dtype,
929
+ )
930
+ layer.workspace = workspace
931
+
932
elif self.fp8_backend in [
933
Fp8MoeBackend.FLASHINFER_CUTLASS,
934
Fp8MoeBackend.FLASHINFER_TRTLLM,
@@ -937,17 +954,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):
954
replace_parameter(layer, f"w13_{self.weight_scale_name}", w13_weight_scale)
955
replace_parameter(layer, f"w2_{self.weight_scale_name}", w2_weight_scale)
956
940
- # TODO(rob): we do this after replace_parameter() because
941
- # prepare_moe_fp8_layer_for_marlin uses on the layer's params
942
- # directly. We will refactor this in a follow up PR.
943
- if self.fp8_backend == Fp8MoeBackend.MARLIN:
944
- prepare_moe_fp8_layer_for_marlin(
945
- layer, False, input_dtype=self.marlin_input_dtype
946
- )
947
- # Activations not quantized for marlin.
948
- del layer.w13_input_scale
949
- del layer.w2_input_scale
950
-
957
def _setup_kernel(self, layer: Module) -> None:
958
"""Setup Modular Kernel for TP Case"""
959
# NOTE(rob): this is a WIP refactor. We are first migrating
@@ -1194,20 +1200,14 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1200
) -> FusedMoEQuantConfig | None:
1201
if self.fp8_backend == Fp8MoeBackend.MARLIN:
1202
return fp8_w8a16_moe_quant_config(
1197
- w1_scale=layer.w13_weight_scale,
1198
- w2_scale=layer.w2_weight_scale,
1203
+ w1_scale=getattr(layer, f"w13_{self.weight_scale_name}"),
1204
+ w2_scale=getattr(layer, f"w2_{self.weight_scale_name}"),
1205
block_shape=self.weight_block_size,
1206
)
1207
1208
return fp8_w8a8_moe_quant_config(
1203
- w1_scale=(
1204
- layer.w13_weight_scale_inv
1205
- if self.block_quant
1206
- else layer.w13_weight_scale
1207
- ),
1208
- w2_scale=(
1209
- layer.w2_weight_scale_inv if self.block_quant else layer.w2_weight_scale
1210
- ),
1209
+ w1_scale=getattr(layer, f"w13_{self.weight_scale_name}"),
1210
+ w2_scale=getattr(layer, f"w2_{self.weight_scale_name}"),
1211
a1_scale=layer.w13_input_scale,
1212
a2_scale=layer.w2_input_scale,
1213
block_shape=self.weight_block_size,
vllm/model_executor/layers/quantization/quark/quark_moe.py
+20
-4
@@ -315,10 +315,26 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
315
layer.w2_weight = torch.nn.Parameter(shuffled_w2, requires_grad=False)
316
317
elif self.use_marlin:
318
- prepare_moe_fp8_layer_for_marlin(layer, False)
319
- # Activations not quantized for marlin.
320
- del layer.w13_input_scale
321
- del layer.w2_input_scale
318
+ (workspace, w13_weight, w2_weight, w13_weight_scale, w2_weight_scale) = (
319
+ prepare_moe_fp8_layer_for_marlin(
320
+ layer,
321
+ layer.w13_weight,
322
+ layer.w2_weight,
323
+ layer.w13_weight_scale,
324
+ layer.w2_weight_scale,
325
+ )
326
+ )
327
+ layer.workspace = workspace
328
+ # TODO(rob): once we apply refactor to Quark, switch to using
329
+ # replace_parameter for compatibility with reloading in RL.
330
+ layer.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False)
331
+ layer.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False)
332
+ layer.w13_weight_scale = torch.nn.Parameter(
333
+ w13_weight_scale, requires_grad=False
334
+ )
335
+ layer.w2_weight_scale = torch.nn.Parameter(
336
+ w2_weight_scale, requires_grad=False
337
+ )
338
339
def get_fused_moe_quant_config(
340
self, layer: torch.nn.Module
vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py
+33
-46
@@ -199,9 +199,18 @@ def prepare_fp8_layer_for_marlin(
199
200
def prepare_moe_fp8_layer_for_marlin(
201
layer: torch.nn.Module,
202
- size_k_first: bool = True,
202
+ w13_weight: torch.Tensor,
203
+ w2_weight: torch.Tensor,
204
+ w13_weight_scale: torch.Tensor,
205
+ w2_weight_scale: torch.Tensor,
206
input_dtype: torch.dtype | None = None,
204
-) -> None:
207
+) -> tuple[
208
+ torch.Tensor, # workspace
209
+ torch.Tensor, # w13_weight
210
+ torch.Tensor, # w2_weight
211
+ torch.Tensor, # w13_weight_scale
212
+ torch.Tensor, # w2_weight_scale
213
+]:
214
logger.warning_once(
215
"Your GPU does not have native support for FP8 computation but "
216
"FP8 quantization is being used. Weight-only FP8 compression will "
@@ -209,7 +218,7 @@ def prepare_moe_fp8_layer_for_marlin(
218
"performance for compute-heavy workloads."
219
)
220
if input_dtype is not None and input_dtype.itemsize == 1:
212
- raise RuntimeError("Marlin W8A8 is not supported.")
221
+ raise NotImplementedError("Marlin W8A8 is not supported.")
222
223
e = layer.num_experts
224
k = layer.hidden_size
@@ -218,53 +227,40 @@ def prepare_moe_fp8_layer_for_marlin(
227
228
# WORKSPACE
229
device = layer.w13_weight.device
221
- layer.workspace = marlin_make_workspace_new(device, 4)
230
+ workspace = marlin_make_workspace_new(device, 4)
231
perm = torch.empty(0, dtype=torch.int, device=device)
232
233
# WEIGHT
234
# Repack weights to marlin format
226
- for name in ["w13_weight", "w2_weight"]:
227
- weight = getattr(layer, name)
235
+ def repack_weight(name: str, weight: torch.Tensor) -> torch.Tensor:
236
tensor_list = []
237
if "w13" in name:
238
size_n, size_k = n * 2, k
239
else:
240
size_n, size_k = k, n
241
234
- if size_k_first:
235
- assert weight.shape == (e, size_k, size_n)
236
- else:
237
- assert weight.shape == (e, size_n, size_k)
242
+ assert weight.shape == (e, size_n, size_k)
243
244
for i in range(e):
240
- qweight = pack_fp8_to_int32(weight[i], size_k_first)
241
- if not size_k_first:
242
- qweight = qweight.T.contiguous()
245
+ qweight = pack_fp8_to_int32(weight[i], size_k_first=False)
246
+ qweight = qweight.T.contiguous()
247
248
marlin_qweight = ops.gptq_marlin_repack(
249
b_q_weight=qweight, perm=perm, size_k=size_k, size_n=size_n, num_bits=8
250
)
251
tensor_list.append(marlin_qweight)
252
249
- weight = torch.cat([x.unsqueeze(0) for x in tensor_list], 0)
250
- weight = torch.nn.Parameter(weight, requires_grad=False)
253
+ return torch.cat([x.unsqueeze(0) for x in tensor_list], 0)
254
252
- setattr(layer, name, weight)
255
+ w13_weight = repack_weight("w13", w13_weight)
256
+ w2_weight = repack_weight("w2", w2_weight)
257
258
# WEIGHT SCALES
259
# Permute scales
260
group_size = -1 if weight_block_size is None else weight_block_size[1]
261
258
- for name in ["w13", "w2"]:
259
- if name + "_weight_scale" in dir(layer):
260
- new_name = name + "_weight_scale"
261
- scales = getattr(layer, new_name).to(layer.orig_dtype)
262
- delattr(layer, new_name)
263
- elif name + "_weight_scale_inv" in dir(layer):
264
- new_name = name + "_weight_scale_inv"
265
- scales = getattr(layer, new_name).to(layer.orig_dtype)
266
- delattr(layer, new_name)
267
-
262
+ def permute_scales(scales: torch.Tensor, name: str) -> torch.Tensor:
263
+ scales = scales.to(layer.orig_dtype)
264
tensor_list = []
265
if "w13" in name:
266
size_n, size_k = n * 2, k
@@ -294,8 +290,7 @@ def prepare_moe_fp8_layer_for_marlin(
290
# block-wise quantization -> group-wise quantization
291
# (e, size_k // block_size[1], ceil(size_n / block_size[0]))
292
# =>(repeat)=> (e, size_k // block_size[1], size_n)
297
- if not size_k_first:
298
- scales = scales.permute(0, 2, 1)
293
+ scales = scales.permute(0, 2, 1)
294
block_n = weight_block_size[0]
295
scales = scales.repeat_interleave(block_n, 2)
296
# size_n may not divisible by block_size[0]
@@ -310,26 +305,18 @@ def prepare_moe_fp8_layer_for_marlin(
305
scales = torch.cat([x.unsqueeze(0) for x in tensor_list], 0)
306
if input_dtype != torch.float8_e4m3fn:
307
scales = fp8_fused_exponent_bias_into_scales(scales)
313
- scales = torch.nn.Parameter(scales, requires_grad=False)
314
-
315
- setattr(layer, name + "_weight_scale", scales)
308
+ return scales
309
317
- # BIAS
318
- # Permute bias
319
- for name in ["w13_bias", "w2_bias"]:
320
- if not hasattr(layer, name):
321
- continue
322
- bias = getattr(layer, name).to(layer.orig_dtype)
310
+ w13_weight_scale = permute_scales(w13_weight_scale, "w13")
311
+ w2_weight_scale = permute_scales(w2_weight_scale, "w2")
312
324
- tensor_list = []
325
- for i in range(e):
326
- expert_bias = bias[i]
327
-
328
- tensor_list.append(marlin_permute_bias(expert_bias))
329
-
330
- bias = torch.cat([x.unsqueeze(0) for x in tensor_list], 0)
331
- bias = torch.nn.Parameter(bias, requires_grad=False)
332
- setattr(layer, name, bias)
313
+ return (
314
+ workspace,
315
+ w13_weight,
316
+ w2_weight,
317
+ w13_weight_scale,
318
+ w2_weight_scale,
319
+ )
320
321
322
def pack_fp8_to_int32(