[lora] enable bias support for fused moe lora

Signed-off-by: AlpinDale <alpindale@gmail.com>

Author: AlpinDale

aphrodite/lora/layers/fused_moe.py

@@ -9,20 +9,17 @@
 from aphrodite.distributed.parallel_state import get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size
 from aphrodite.lora.layers.base import BaseLayerWithLoRA
 from aphrodite.modeling.layers.fused_moe import FusedMoE
-from aphrodite.modeling.layers.fused_moe.config import (
-    FUSED_MOE_UNQUANTIZED_CONFIG,
-    _get_config_dtype_str,
-    mxfp4_w4a16_moe_quant_config,
-)
+from aphrodite.modeling.layers.fused_moe.config import _get_config_dtype_str
 from aphrodite.modeling.layers.fused_moe.fused_marlin_moe import modular_marlin_fused_moe
 from aphrodite.modeling.layers.fused_moe.fused_moe import modular_triton_fused_moe, try_get_optimal_moe_config
-from aphrodite.quantization.mxfp4 import Mxfp4Config
 
 
 class FusedMoEWithLoRA(BaseLayerWithLoRA):
     def __init__(self, base_layer: FusedMoE) -> None:
         super().__init__()
         self.base_layer = base_layer
+
+        assert not self.base_layer.use_ep, "EP support for Fused MoE LoRA is not implemented yet."
         self.tp_size = get_tensor_model_parallel_world_size()
         self.tp_rank = get_tensor_model_parallel_rank()
         self.device = base_layer.w2_weight.device
@@ -32,17 +29,8 @@ def _inject_lora_into_fused_moe(self):
         moe_state_dict = {}
         top_k = self.base_layer.top_k
 
-        if self.base_layer.quant_config is None:
-            quant_config = FUSED_MOE_UNQUANTIZED_CONFIG
-        elif not isinstance(self.base_layer.quant_config, Mxfp4Config):
-            quant_config = self.base_layer.quant_config
-        else:
-            quant_config = mxfp4_w4a16_moe_quant_config(
-                w1_bias=self.base_layer.w13_bias,
-                w2_bias=self.base_layer.w2_bias,
-                w1_scale=self.base_layer.w13_weight_scale,
-                w2_scale=self.base_layer.w2_weight_scale,
-            )
+        self.base_layer.ensure_moe_quant_config_init()
+        quant_config = self.base_layer.quant_method.moe_quant_config
 
         m_fused_moe_fn = (
             modular_triton_fused_moe(quant_config, shared_experts=self.base_layer.shared_experts)
@@ -70,7 +58,6 @@ def wrapper(*args, **kwargs):
                 hidden_states = moe_state_dict["hidden_states"]
                 topk_weights = moe_state_dict["topk_weights"]
                 curr_topk_ids = moe_state_dict["topk_ids"]
-                global_num_experts = moe_state_dict["global_num_experts"]
                 expert_map = moe_state_dict["expert_map"]
 
                 config_dtype = _get_config_dtype_str(
@@ -102,7 +89,7 @@ def wrapper(*args, **kwargs):
                     curr_topk_ids,
                     num_tokens,
                     config["BLOCK_SIZE_M"],
-                    global_num_experts,
+                    self.base_layer.local_num_experts,
                     max_loras,
                     expert_map,
                 )
@@ -210,12 +197,10 @@ def create_lora_weights(
     ) -> None:
         """Initializes lora matrices."""
 
-        assert not self.base_layer.use_ep, "EP support for Fused MoE LoRA is not implemented yet."
-
         self.w1_lora_a_stacked = torch.zeros(
             (
                 max_loras,
-                self.base_layer.global_num_experts,
+                self.base_layer.local_num_experts,
                 lora_config.max_lora_rank,
                 self.base_layer.hidden_size,
             ),
@@ -225,7 +210,7 @@ def create_lora_weights(
         self.w1_lora_b_stacked = torch.zeros(
             (
                 max_loras,
-                self.base_layer.global_num_experts,
+                self.base_layer.local_num_experts,
                 self.base_layer.intermediate_size_per_partition,
                 lora_config.max_lora_rank,
             ),
@@ -236,7 +221,7 @@ def create_lora_weights(
         self.w2_lora_a_stacked = torch.zeros(
             (
                 max_loras,
-                self.base_layer.global_num_experts,
+                self.base_layer.local_num_experts,
                 lora_config.max_lora_rank,
                 self.base_layer.intermediate_size_per_partition,
             ),
@@ -246,7 +231,7 @@ def create_lora_weights(
         self.w2_lora_b_stacked = torch.zeros(
             (
                 max_loras,
-                self.base_layer.global_num_experts,
+                self.base_layer.local_num_experts,
                 self.base_layer.hidden_size,
                 lora_config.max_lora_rank,
             ),
@@ -257,7 +242,7 @@ def create_lora_weights(
         self.w3_lora_a_stacked = torch.zeros(
             (
                 max_loras,
-                self.base_layer.global_num_experts,
+                self.base_layer.local_num_experts,
                 lora_config.max_lora_rank,
                 self.base_layer.hidden_size,
             ),
@@ -267,7 +252,7 @@ def create_lora_weights(
         self.w3_lora_b_stacked = torch.zeros(
             (
                 max_loras,
-                self.base_layer.global_num_experts,
+                self.base_layer.local_num_experts,
                 self.base_layer.intermediate_size_per_partition,
                 lora_config.max_lora_rank,
             ),
@@ -280,7 +265,7 @@ def create_lora_weights(
         self.lora_a_stacked = []
         self.lora_b_stacked = []
         for lora_id in range(max_loras):
-            for experts_id in range(self.base_layer.global_num_experts):
+            for experts_id in range(self.base_layer.local_num_experts):
                 # gate_proj,down_proj,up_proj
                 self.lora_a_stacked.append(self.w1_lora_a_stacked[lora_id][experts_id])
                 self.lora_a_stacked.append(self.w2_lora_a_stacked[lora_id][experts_id])

aphrodite/modeling/layers/fused_moe/layer.py

@@ -629,8 +629,6 @@ def forward_cuda(
                 apply_router_weight_on_input=apply_router_weight_on_input,
             )
         elif self.fused_experts is not None:
-            if self.moe.has_bias:
-                raise ValueError("FusedMoEModularKernel does not support bias.")
             result = self.fused_experts(
                 hidden_states=x,
                 w1=layer.w13_weight,