feat: mega moe weight in place

Author: niushengxiao

lightllm/common/basemodel/layer_weights/base_layer_weight.py

@@ -38,6 +38,7 @@ def verify_load(self):
                 else:
                     layer_num = None
                 assert attr.verify_load(), f"Loading {attr_name} of layers {layer_num} fails."
+                attr.finalize_load()
 
     def _cuda(self, cpu_tensor):
         return cpu_tensor.contiguous().to(self.data_type_).cuda(get_current_device_id())

lightllm/common/basemodel/layer_weights/meta_weights/base_weight.py

@@ -21,6 +21,9 @@ def _create_weight(self):
     def verify_load(self) -> bool:
         pass
 
+    def finalize_load(self) -> None:
+        pass
+
 
 class BaseWeightTpl(BaseWeight):
     def __init__(self, tp_rank: int = None, tp_world_size: int = None, data_type: torch.dtype = None):

lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py

@@ -284,6 +284,16 @@ def verify_load(self):
         )
         return weight_load_ok and per_expert_scale_load_ok and e_score_correction_bias_load_ok
 
+    def finalize_load(self):
+        if self.enable_ep_moe:
+            from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep import (
+                transform_mega_moe_weights_in_place,
+                use_sm100_mega_moe,
+            )
+
+            if use_sm100_mega_moe(self.quant_method):
+                transform_mega_moe_weights_in_place(self.w13, self.w2)
+
     def _create_weight(self):
         intermediate_size = self.split_inter_size
         self.e_score_correction_bias = None

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

@@ -22,7 +22,7 @@
 from lightllm.utils.device_utils import is_sm100_gpu
 
 logger = init_logger(__name__)
-_MEGA_MOE_STATES: Dict[Tuple[int, int, int, int], Dict[str, Any]] = {}
+_MEGA_MOE_STATS: Dict[Tuple[int, int, int, int], torch.Tensor] = {}
 SUPPORTED_EP_EXPERT_DTYPES = ("deepgemm-fp8w8a8-b128", "deepgemm-fp4fp8-b32")
 
 try:
@@ -232,17 +232,21 @@ def masked_group_gemm(
     return gemm_out_b
 
 
-def _get_mega_moe_cache_state(w13: Any, w2: Any):
+def _get_mega_moe_cumulative_stats(w13: Any, w2: Any):
     state_key = (
         w13.weight.data_ptr(),
         w13.weight_scale.data_ptr(),
         w2.weight.data_ptr(),
         w2.weight_scale.data_ptr(),
     )
-    return _MEGA_MOE_STATES.setdefault(state_key, {})
+    stats = _MEGA_MOE_STATS.get(state_key)
+    if stats is None:
+        stats = torch.zeros((w13.weight.shape[0],), device=w13.weight.device, dtype=torch.int32)
+        _MEGA_MOE_STATS[state_key] = stats
+    return stats
 
 
-def _transform_mega_moe_weights_in_place(w13: Any, w2: Any):
+def transform_mega_moe_weights_in_place(w13: Any, w2: Any):
     """Convert to Mega MoE layout without retaining a second weight copy."""
     transformed_l1, transformed_l2 = deep_gemm.transform_weights_for_mega_moe(
         (w13.weight, w13.weight_scale),
@@ -251,21 +255,6 @@ def _transform_mega_moe_weights_in_place(w13: Any, w2: Any):
     w13.weight.copy_(transformed_l1[0])
     w13.weight_scale.copy_(transformed_l1[1])
     w2.weight_scale.copy_(transformed_l2[1])
-    return (w13.weight, w13.weight_scale), (w2.weight, w2.weight_scale)
-
-
-def _get_mega_moe_weights(w13: Any, w2: Any, state: Dict[str, Any]):
-    if "weight_cache" not in state:
-        state["weight_cache"] = _transform_mega_moe_weights_in_place(w13, w2)
-    return state["weight_cache"]
-
-
-def _get_mega_moe_cumulative_stats(num_local_experts: int, device: torch.device, state: Dict[str, Any]):
-    stats = state.get("stats")
-    if stats is None or stats.numel() != num_local_experts or stats.device != device:
-        stats = torch.zeros((num_local_experts,), device=device, dtype=torch.int32)
-        state["stats"] = stats
-    return stats
 
 
 def mega_moe_impl(
@@ -289,9 +278,9 @@ def mega_moe_impl(
             f"Mega MoE got {num_tokens} tokens, exceeding num_max_tokens_per_rank={buffer.num_max_tokens_per_rank}"
         )
 
-    state = _get_mega_moe_cache_state(w13, w2)
-    l1_weights, l2_weights = _get_mega_moe_weights(w13, w2, state)
-    stats = _get_mega_moe_cumulative_stats(w13.weight.shape[0], hidden_states.device, state)
+    l1_weights = (w13.weight, w13.weight_scale)
+    l2_weights = (w2.weight, w2.weight_scale)
+    stats = _get_mega_moe_cumulative_stats(w13, w2)
     _prepare_mega_moe_buffer(hidden_states, topk_ids, topk_weights, buffer, quant_method.block_size)
 
     output = torch.empty_like(hidden_states)