feat: add --expert_dtype param

Author: niushengxiao

docs/CN/source/tutorial/api_server_args.rst

@@ -464,6 +464,14 @@ PD 分离模式参数
 
     示例可以在 test/advanced_config/mixed_quantization/llamacls-mix-down.yaml 中找到。
 
+.. option:: --expert_dtype
+
+    EP MoE 专家量化类型，可选值：
+
+    * ``deepgemm-fp8w8a8-b128``
+    * ``deepgemm-fp4fp8-b32``，仅支持 SM100 GPU
+    * ``None`` (默认)
+
 .. option:: --vit_quant_type
 
     ViT 量化方法，可选值：

docs/EN/source/tutorial/api_server_args.rst

@@ -465,6 +465,14 @@ Quantization Parameters
 
     Examples can be found in test/advanced_config/mixed_quantization/llamacls-mix-down.yaml.
 
+.. option:: --expert_dtype
+
+    Expert quantization dtype for EP MoE, optional values:
+
+    * ``deepgemm-fp8w8a8-b128``
+    * ``deepgemm-fp4fp8-b32``: SM100 GPUs only
+    * ``None`` (default)
+
 .. option:: --vit_quant_type
 
     ViT quantization method, optional values:

lightllm/common/basemodel/basemodel.py

@@ -85,6 +85,7 @@ def __init__(self, kvargs):
         self.disable_cudagraph = kvargs.get("disable_cudagraph", False)
         self.quant_type = kvargs.get("quant_type", "none")
         self.quant_cfg_path = kvargs.get("quant_cfg", None)
+        self.expert_dtype = kvargs.get("expert_dtype", None)
         self.mem_fraction = kvargs.get("mem_fraction", 0.9)
         self.tp_world_size_ = get_dp_world_size()
         self.enable_tpsp_mix_mode = get_env_start_args().enable_tpsp_mix_mode
@@ -156,7 +157,7 @@ def _verify_params(self):
         return
 
     def _init_quant(self):
-        self.quant_cfg = Quantcfg(self.config, self.quant_type, self.quant_cfg_path)
+        self.quant_cfg = Quantcfg(self.config, self.quant_type, self.quant_cfg_path, self.expert_dtype)
         logger.info(f"Initial quantization. " f"The default quantization method is {self.quant_cfg.quant_type}")
 
     def _init_weights(self, start_layer_index=0):

lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py

@@ -13,7 +13,7 @@
     get_ep_num_sms,
     masked_group_gemm,
     deepgemm_grouped_fp8_nt_contiguous,
-    use_sm100_fp4_moe,
+    use_sm100_mega_moe,
 )
 from lightllm.common.basemodel.triton_kernel.quantization.fp8act_quant_kernel import (
     per_token_group_quant_fp8,
@@ -153,7 +153,7 @@ def select_experts_and_quant_input(
             scoring_func=scoring_func,
         )
         w13_weight, w13_scale = w13.weight, w13.weight_scale
-        if use_sm100_fp4_moe(self.quant_method):
+        if use_sm100_mega_moe(self.quant_method):
             from deep_gemm.utils import per_token_cast_to_fp8
 
             qinput_tensor = per_token_cast_to_fp8(

lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py

@@ -22,6 +22,7 @@
 
 logger = init_logger(__name__)
 _MEGA_MOE_STATES: Dict[Tuple[int, int, int, int], Dict[str, Any]] = {}
+SUPPORTED_EP_EXPERT_DTYPES = ("deepgemm-fp8w8a8-b128", "deepgemm-fp4fp8-b32")
 
 try:
     from deep_ep import Buffer, EventOverlap
@@ -37,10 +38,27 @@ def get_ep_num_sms() -> int:
     return getattr(dist_group_manager, "ep_num_sms", None) or 0
 
 
-def use_sm100_fp4_moe(quant_method: Any) -> bool:
+def use_sm100_mega_moe(quant_method: Any) -> bool:
     return is_sm100_gpu() and quant_method.method_name == "deepgemm-fp4fp8-b32"
 
 
+def check_ep_expert_dtype(quant_method: Any):
+    expert_dtype = getattr(quant_method, "method_name", None)
+    if expert_dtype not in SUPPORTED_EP_EXPERT_DTYPES:
+        raise ValueError(
+            "EP MoE requires --expert_dtype to be one of "
+            f"{list(SUPPORTED_EP_EXPERT_DTYPES)}, but got `{expert_dtype}`. "
+            "Please start with --expert_dtype deepgemm-fp8w8a8-b128 or "
+            "--expert_dtype deepgemm-fp4fp8-b32. Note that deepgemm-fp4fp8-b32 "
+            "is only supported on SM100 GPUs."
+        )
+    if expert_dtype == "deepgemm-fp4fp8-b32" and not is_sm100_gpu():
+        raise RuntimeError(
+            "--expert_dtype deepgemm-fp4fp8-b32 requires an SM100 GPU for EP MoE; "
+            "please use --expert_dtype deepgemm-fp8w8a8-b128 on non-SM100 GPUs."
+        )
+
+
 def masked_group_gemm(
     recv_x: Tuple[torch.Tensor, torch.Tensor],
     masked_m: torch.Tensor,
@@ -155,10 +173,10 @@ def do_fused_experts(
     is_prefill: Optional[bool],
     previous_event: Optional[Any] = None,
 ):
-    if use_sm100_fp4_moe(quant_method):
+    check_ep_expert_dtype(quant_method)
+    if use_sm100_mega_moe(quant_method):
         return mega_moe_impl(hidden_states, w13, w2, topk_weights, topk_idx, quant_method)
 
-    use_fp8_w8a8 = quant_method.method_name != "none"
     buffer = dist_group_manager.ep_buffer if is_prefill else dist_group_manager.ep_low_latency_buffer
     return fused_experts_impl(
         hidden_states=hidden_states,
@@ -169,8 +187,8 @@ def do_fused_experts(
         num_experts=num_experts,
         buffer=buffer,
         is_prefill=is_prefill,
-        use_fp8_w8a8=use_fp8_w8a8,
-        use_fp8_all2all=use_fp8_w8a8,
+        use_fp8_w8a8=True,
+        use_fp8_all2all=True,
         use_int8_w8a16=False,
         w1_scale=w13.weight_scale,
         w2_scale=w2.weight_scale,

lightllm/common/quantization/__init__.py

@@ -7,17 +7,36 @@
 from .awq import *
 from .no_quant import *
 from lightllm.utils.log_utils import init_logger
+from lightllm.utils.device_utils import is_sm100_gpu
 
 logger = init_logger(__name__)
 
+SUPPORTED_EXPERT_DTYPES = ("deepgemm-fp8w8a8-b128", "deepgemm-fp4fp8-b32")
+
 
 class Quantcfg:
-    def __init__(self, network_config, quant_type="none", custom_cfg_path=None):
+    def __init__(self, network_config, quant_type="none", custom_cfg_path=None, expert_dtype=None):
         self.layer_num = network_config["n_layer"]
         self.quant_type = quant_type
         self.network_config_ = network_config
         self._parse_custom_cfg(custom_cfg_path)
         self._parse_network_config(network_config)
+        self._apply_custom_expert_dtype(expert_dtype)
+
+    def _apply_custom_expert_dtype(self, expert_dtype):
+        if expert_dtype is None:
+            return
+        if expert_dtype not in SUPPORTED_EXPERT_DTYPES:
+            raise ValueError(
+                f"unsupported --expert_dtype `{expert_dtype}`; expected one of {list(SUPPORTED_EXPERT_DTYPES)}"
+            )
+        if not is_sm100_gpu() and expert_dtype == "deepgemm-fp4fp8-b32":
+            raise RuntimeError(
+                f"deepgemm-fp4fp8-b32 requires an SM100 GPU; " "please use deepgemm-fp8w8a8-b128 on non-SM100 GPUs."
+            )
+        for layer_num in range(self.layer_num):
+            self.quant_cfg[layer_num]["fused_moe"] = expert_dtype
+        logger.info(f"select fused_moe quant way from --expert_dtype: {expert_dtype}")
 
     def _parse_network_config(self, network_config):
         hf_quantization_config = network_config.get("quantization_config", None)

lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py

@@ -7,6 +7,7 @@
 from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer
 from lightllm.models.deepseek2.triton_kernel.rotary_emb import rotary_emb_fwd
 from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo
+from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep import use_sm100_mega_moe
 from functools import partial
 from lightllm.models.llama.yarn_rotary_utils import get_deepseek_mscale
 from lightllm.utils.envs_utils import get_env_start_args
@@ -295,7 +296,7 @@ def overlap_tpsp_token_forward(
         infer_state1: Deepseek2InferStateInfo,
         layer_weight: Deepseek2TransformerLayerWeight,
     ):
-        if not self.is_moe or layer_weight.experts.use_sm100_mega_moe():
+        if not self.is_moe or use_sm100_mega_moe(layer_weight.experts.quant_method):
             return super().overlap_tpsp_token_forward(
                 input_embdings, input_embdings1, infer_state, infer_state1, layer_weight
             )
@@ -421,7 +422,7 @@ def overlap_tpsp_context_forward(
         infer_state1: Deepseek2InferStateInfo,
         layer_weight: Deepseek2TransformerLayerWeight,
     ):
-        if not self.is_moe or layer_weight.experts.use_sm100_mega_moe():
+        if not self.is_moe or use_sm100_mega_moe(layer_weight.experts.quant_method):
             return super().overlap_tpsp_context_forward(
                 input_embdings, input_embdings1, infer_state, infer_state1, layer_weight
             )

lightllm/models/qwen3_moe/layer_infer/transformer_layer_infer.py

@@ -6,6 +6,7 @@
 from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer
 from lightllm.models.llama.infer_struct import LlamaInferStateInfo
 from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd
+from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep import use_sm100_mega_moe
 from lightllm.utils.dist_utils import get_global_world_size
 from lightllm.utils.envs_utils import get_env_start_args
 
@@ -133,7 +134,7 @@ def overlap_tpsp_token_forward(
         infer_state1: LlamaInferStateInfo,
         layer_weight: Qwen3MOETransformerLayerWeight,
     ):
-        if not self.is_moe or layer_weight.experts.use_sm100_mega_moe():
+        if not self.is_moe or use_sm100_mega_moe(layer_weight.experts.quant_method):
             return super().overlap_tpsp_token_forward(
                 input_embdings, input_embdings1, infer_state, infer_state1, layer_weight
             )
@@ -245,7 +246,7 @@ def overlap_tpsp_context_forward(
         infer_state1: LlamaInferStateInfo,
         layer_weight: Qwen3MOETransformerLayerWeight,
     ):
-        if not self.is_moe or layer_weight.experts.use_sm100_mega_moe():
+        if not self.is_moe or use_sm100_mega_moe(layer_weight.experts.quant_method):
             return super().overlap_tpsp_context_forward(
                 input_embdings, input_embdings1, infer_state, infer_state1, layer_weight
             )

lightllm/server/api_cli.py

@@ -620,6 +620,15 @@ def make_argument_parser() -> argparse.ArgumentParser:
         help="""Path of quantization config. It can be used for mixed quantization.
             Examples can be found in test/advanced_config/mixed_quantization/llamacls-mix-down.yaml.""",
     )
+    parser.add_argument(
+        "--expert_dtype",
+        type=str,
+        default=None,
+        choices=["deepgemm-fp8w8a8-b128", "deepgemm-fp4fp8-b32"],
+        help="""Expert quantization dtype for EP MoE. Supported values are
+            deepgemm-fp8w8a8-b128 and deepgemm-fp4fp8-b32. Note that
+            deepgemm-fp4fp8-b32 is only supported on SM100 GPUs.""",
+    )
     parser.add_argument(
         "--vit_quant_type",
         type=str,

lightllm/server/core/objs/start_args_type.py

@@ -133,6 +133,9 @@ class StartArgs:
     graph_max_len_in_batch: int = field(default=0)
     quant_type: Optional[str] = field(default=None)
     quant_cfg: Optional[str] = field(default=None)
+    expert_dtype: Optional[str] = field(
+        default=None, metadata={"choices": ["deepgemm-fp8w8a8-b128", "deepgemm-fp4fp8-b32"]}
+    )
     vit_quant_type: Optional[str] = field(default=None)
     vit_quant_cfg: Optional[str] = field(default=None)
     llm_prefill_att_backend: List[str] = field(