@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / b471092d3a

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 11, 03:45 UTC

[MoE Refactor][4/N] Marlin Fp8 Mk (#31036)

Robert Shaw committed Dec 21, 2025 at 12:37 UTC b471092d3a983485243dd153777d026bf5f49183

4 files changed +85 -63

tests/quantization/test_fp8.py

@@ -15,6 +15,7 @@ from vllm.model_executor.layers.quantization.fp8 import (

15 Fp8Config,

16 Fp8KVCacheMethod,

17 Fp8LinearMethod,

18 + Fp8MoeBackend,

19 Fp8MoEMethod,

20 )

21 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

@@ -324,7 +325,10 @@ def test_fp8_reloading(

325 weight_loader=default_weight_loader,

326 )

327

328 + # Fp8LinearMethod uses use_marlin

329 + # Fp8MoEMethod uses fp8_backend

330 method.use_marlin = use_marlin

331 + method.fp8_backend = Fp8MoeBackend.MARLIN if use_marlin else None

332

333 # capture weights format during loading

334 original_metadata = [

vllm/model_executor/layers/fused_moe/config.py

+30 -2

@@ -19,6 +19,7 @@ from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (

19 OCP_MX_Scheme,

20 )

21 from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape

22 +from vllm.platforms import current_platform

23 from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe

24 from vllm.utils.import_utils import has_triton_kernels

25 from vllm.utils.math_utils import cdiv

@@ -39,6 +40,7 @@ if has_triton_kernels():

40 def _get_config_dtype_str(

41 dtype: torch.dtype,

42 use_fp8_w8a8: bool = False,

43 + use_fp8_w8a16: bool = False,

44 use_int8_w8a16: bool = False,

45 use_int4_w4a16: bool = False,

46 ocp_mx_scheme: str | None = None,

@@ -50,6 +52,8 @@ def _get_config_dtype_str(

52 """

53 if use_fp8_w8a8:

54 return "fp8_w8a8"

55 + elif use_fp8_w8a16:

56 + return "fp8_w8a16"

57 elif use_int8_w8a16:

58 return "int8_w8a16"

59 elif use_int4_w4a16:

@@ -319,6 +323,10 @@ class FusedMoEQuantConfig:

323 def use_int8_w8a16(self) -> bool:

324 return self._a1.dtype is None and self._w1.dtype == torch.int8

325

326 + @property

327 + def use_fp8_w8a16(self) -> bool:

328 + return self._a1.dtype is None and self._w1.dtype == current_platform.fp8_dtype()

329 +

330 @property

331 def use_int4_w4a16(self) -> bool:

332 return self._a1.dtype is None and self._w1.dtype == "int4"

@@ -362,6 +370,7 @@ class FusedMoEQuantConfig:

370 """

371 return _get_config_dtype_str(

372 use_fp8_w8a8=self.use_fp8_w8a8,

373 + use_fp8_w8a16=self.use_fp8_w8a16,

374 use_int8_w8a16=self.use_int8_w8a16,

375 use_int4_w4a16=self.use_int4_w4a16,

376 ocp_mx_scheme=self.ocp_mx_scheme,

@@ -680,7 +689,6 @@ def int4_w4a16_moe_quant_config(

689 ) -> FusedMoEQuantConfig:

690 """

691 Construct a quant config for 16-bit float activations and int4 weights.

683 - Note: Activations are pre-quantized.

692 """

693 group_shape = GroupShape(*block_shape) if block_shape is not None else None

694 return FusedMoEQuantConfig(

@@ -691,6 +699,27 @@ def int4_w4a16_moe_quant_config(

699 )

700

701

702 +def fp8_w8a16_moe_quant_config(

703 + w1_scale: torch.Tensor,

704 + w2_scale: torch.Tensor,

705 + block_shape: list[int] | None = None,

706 +) -> FusedMoEQuantConfig:

707 + """

708 + Construct a quant config for 16-bit float activations and fp8 weights.

709 + """

710 + group_shape = GroupShape(*block_shape) if block_shape is not None else None

711 + return FusedMoEQuantConfig(

712 + _a1=FusedMoEQuantDesc(),

713 + _a2=FusedMoEQuantDesc(),

714 + _w1=FusedMoEQuantDesc(

715 + current_platform.fp8_dtype(), group_shape, w1_scale, None, None

716 + ),

717 + _w2=FusedMoEQuantDesc(

718 + current_platform.fp8_dtype(), group_shape, w2_scale, None, None

719 + ),

720 + )

721 +

722 +

723 def int8_w8a16_moe_quant_config(

724 w1_scale: torch.Tensor,

725 w2_scale: torch.Tensor,

@@ -700,7 +729,6 @@ def int8_w8a16_moe_quant_config(

729 ) -> FusedMoEQuantConfig:

730 """

731 Construct a quant config for 16-bit float activations and int8 weights.

703 - Note: Activations are pre-quantized.

732 """

733 group_shape = GroupShape(*block_shape) if block_shape is not None else None

734 return FusedMoEQuantConfig(

vllm/model_executor/layers/fused_moe/fused_marlin_moe.py

+17 -21

@@ -13,9 +13,6 @@ from vllm.model_executor.layers.fused_moe.moe_align_block_size import (

13 batched_moe_align_block_size,

14 moe_align_block_size,

15 )

16 -from vllm.model_executor.layers.fused_moe.prepare_finalize import (

17 - MoEPrepareAndFinalizeNoEP,

18 -)

16 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

17 TopKWeightAndReduceDelegate,

18 TopKWeightAndReduceNoOP,

@@ -26,6 +23,7 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (

23 marlin_moe_intermediate_size,

24 marlin_quant_input,

25 )

26 +from vllm.platforms import current_platform

27 from vllm.scalar_type import ScalarType, scalar_types

@@ -542,9 +540,11 @@ class MarlinExpertsBase(mk.FusedMoEPermuteExpertsUnpermute):

540 is_k_full: bool = True,

541 ):

542 # TODO (varun) : Enable activation quantization

545 - assert quant_config.use_mxfp4_w4a16 or quant_config.use_int4_w4a16, (

546 - "Supports only mxfp4_w4a16 or int4_w4a16"

547 - )

543 + assert (

544 + quant_config.use_mxfp4_w4a16

545 + or quant_config.use_int4_w4a16

546 + or quant_config.use_fp8_w8a16

547 + ), "Supports only mxfp4_w4a16, int4_w4a16 or fp8_w8a16"

548 self.w13_g_idx = w13_g_idx

549 self.w2_g_idx = w2_g_idx

550 self.w13_g_idx_sort_indices = w13_g_idx_sort_indices

@@ -555,11 +555,17 @@ class MarlinExpertsBase(mk.FusedMoEPermuteExpertsUnpermute):

555 @property

556 def quant_type_id(self) -> int:

557 # uint4b8 will be set for int4 weight and float4_e2m1f will be used for mxfp4

558 - return (

559 - scalar_types.uint4b8.id

560 - if self.quant_config.use_int4_w4a16

561 - else scalar_types.float4_e2m1f.id

562 - )

558 + if self.quant_config.use_int4_w4a16:

559 + return scalar_types.uint4b8.id

560 + elif self.quant_config.use_mxfp4_w4a16:

561 + return scalar_types.float4_e2m1f.id

562 + elif (

563 + self.quant_config.use_fp8_w8a16

564 + and current_platform.fp8_dtype() == torch.float8_e4m3fn

565 + ):

566 + return scalar_types.float8_e4m3fn.id

567 + else:

568 + raise NotImplementedError("Unsupported quantization type.")

569

570 def moe_problem_size(

571 self,

@@ -711,16 +717,6 @@ class MarlinExperts(MarlinExpertsBase):

717 ops.moe_sum(input, output)

718

719

714 -def modular_marlin_fused_moe(

715 - quant_config: FusedMoEQuantConfig, shared_experts: torch.nn.Module | None = None

716 -) -> mk.FusedMoEModularKernel:

717 - return mk.FusedMoEModularKernel(

718 - MoEPrepareAndFinalizeNoEP(),

719 - MarlinExperts(quant_config),

720 - shared_experts,

721 - )

722 -

723 -

720 class BatchedMarlinExperts(MarlinExpertsBase):

721 def __init__(

722 self,

vllm/model_executor/layers/quantization/fp8.py

+34 -40

@@ -32,8 +32,8 @@ from vllm.model_executor.layers.fused_moe.config import (

32 FusedMoEQuantConfig,

33 RoutingMethodType,

34 fp8_w8a8_moe_quant_config,

35 + fp8_w8a16_moe_quant_config,

36 )

36 -from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_moe

37 from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod

38 from vllm.model_executor.layers.linear import (

39 LinearBase,

@@ -95,7 +95,6 @@ from vllm.model_executor.parameter import (

95 )

96 from vllm.model_executor.utils import replace_parameter, set_weight_attrs

97 from vllm.platforms import current_platform

98 -from vllm.scalar_type import scalar_types

98 from vllm.utils.deep_gemm import (

99 is_deep_gemm_e8m0_used,

100 is_deep_gemm_supported,

@@ -729,7 +728,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):

728 )

729

730 self.marlin_input_dtype = None

732 - self.use_marlin = self.fp8_backend == Fp8MoeBackend.MARLIN

731 self.flashinfer_moe_backend: FlashinferMoeBackend | None = None

732 if self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM:

733 self.flashinfer_moe_backend = FlashinferMoeBackend.TENSORRT_LLM

@@ -1048,7 +1046,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1046 rotate_flashinfer_fp8_moe_weights(w13_weight, w2_weight)

1047 layer.w13_weight.data = w13_weight.data

1048

1051 - if self.use_marlin:

1049 + if self.fp8_backend == Fp8MoeBackend.MARLIN:

1050 prepare_moe_fp8_layer_for_marlin(

1051 layer, False, input_dtype=self.marlin_input_dtype

1052 )

@@ -1091,10 +1089,17 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1089 )

1090 self.use_inplace = False

1091

1094 - elif self.fp8_backend in [Fp8MoeBackend.DEEPGEMM, Fp8MoeBackend.TRITON]:

1092 + elif self.fp8_backend in [

1093 + Fp8MoeBackend.DEEPGEMM,

1094 + Fp8MoeBackend.TRITON,

1095 + Fp8MoeBackend.MARLIN,

1096 + ]:

1097 from vllm.model_executor.layers.fused_moe import (

1098 TritonOrDeepGemmExperts,

1099 )

1100 + from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (

1101 + MarlinExperts,

1102 + )

1103 from vllm.model_executor.layers.fused_moe.prepare_finalize import (

1104 MoEPrepareAndFinalizeNoEP,

1105 )

@@ -1102,12 +1107,19 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1107 config = self.get_fused_moe_quant_config(layer)

1108 assert config is not None

1109 self.moe_quant_config = config

1105 - self.kernel = mk.FusedMoEModularKernel(

1106 - MoEPrepareAndFinalizeNoEP(),

1107 - TritonOrDeepGemmExperts(

1110 + use_marlin = self.fp8_backend == Fp8MoeBackend.MARLIN

1111 + allow_deep_gemm = self.fp8_backend == Fp8MoeBackend.DEEPGEMM

1112 + moe_kernel = (

1113 + MarlinExperts(quant_config=self.moe_quant_config)

1114 + if use_marlin

1115 + else TritonOrDeepGemmExperts(

1116 quant_config=self.moe_quant_config,

1109 - allow_deep_gemm=(self.fp8_backend == Fp8MoeBackend.DEEPGEMM),

1110 - ),

1117 + allow_deep_gemm=allow_deep_gemm,

1118 + )

1119 + )

1120 +

1121 + self.kernel = mk.FusedMoEModularKernel(

1122 + MoEPrepareAndFinalizeNoEP(), moe_kernel

1123 )

1124 self.use_inplace = True

1125

@@ -1116,9 +1128,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1128 routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,

1129 ) -> mk.FusedMoEPrepareAndFinalize | None:

1130 if (

1119 - current_platform.is_xpu()

1120 - or self.rocm_aiter_moe_enabled

1121 - or self.use_marlin

1131 + self.rocm_aiter_moe_enabled

1132 + or self.fp8_backend == Fp8MoeBackend.MARLIN

1133 or self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM

1134 ):

1135 return None

@@ -1150,7 +1161,9 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1161 TritonOrDeepGemmExperts,

1162 )

1163

1153 - assert not self.use_marlin and not self.rocm_aiter_moe_enabled, (

1164 + assert (

1165 + self.fp8_backend != Fp8MoeBackend.MARLIN

1166 + ) and not self.rocm_aiter_moe_enabled, (

1167 "Marlin and ROCm AITER are not supported with all2all yet."

1168 )

1169

@@ -1207,8 +1220,12 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1220 def get_fused_moe_quant_config(

1221 self, layer: torch.nn.Module

1222 ) -> FusedMoEQuantConfig | None:

1210 - if self.use_marlin:

1211 - return None

1223 + if self.fp8_backend == Fp8MoeBackend.MARLIN:

1224 + return fp8_w8a16_moe_quant_config(

1225 + w1_scale=layer.w13_weight_scale,

1226 + w2_scale=layer.w2_weight_scale,

1227 + block_shape=self.weight_block_size,

1228 + )

1229

1230 return fp8_w8a8_moe_quant_config(

1231 w1_scale=(

@@ -1314,29 +1331,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1331 expert_map=layer.expert_map,

1332 quant_config=self.moe_quant_config,

1333 )

1317 - elif self.use_marlin:

1318 - # TODO(rob): convert this to MK.

1319 - assert layer.activation == "silu", (

1320 - f"{layer.activation} not supported for Marlin MoE."

1321 - )

1322 - result = fused_marlin_moe(

1323 - x,

1324 - layer.w13_weight,

1325 - layer.w2_weight,

1326 - None,

1327 - None,

1328 - layer.w13_weight_scale,

1329 - layer.w2_weight_scale,

1330 - router_logits,

1331 - topk_weights,

1332 - topk_ids,

1333 - quant_type_id=scalar_types.float8_e4m3fn.id,

1334 - apply_router_weight_on_input=layer.apply_router_weight_on_input,

1335 - global_num_experts=layer.global_num_experts,

1336 - expert_map=layer.expert_map,

1337 - input_dtype=self.marlin_input_dtype,

1338 - workspace=layer.workspace,

1339 - )

1334 else:

1335 result = self.kernel(

1336 x,

@@ -1495,7 +1489,7 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod):

1489 replace_parameter(layer, "w2_weight", shuffled_w2)

1490

1491 # Rushuffle weights for MARLIN if needed.

1498 - if self.use_marlin:

1492 + if self.fp8_backend == Fp8MoeBackend.MARLIN:

1493 prepare_moe_fp8_layer_for_marlin(

1494 layer, False, input_dtype=self.marlin_input_dtype

1495 )