MoE intermediate_scale stack

Author: yanfeich

fastdeploy/model_executor/layers/backends/intel_hpu/moe/fused_moe_hpu_backend.py

@@ -59,18 +59,17 @@ def init_ep(self, layer: nn.Layer) -> None:
         """
         return
 
-    def apply(
+    def apply_tp(
         self,
         layer: nn.Layer,
         x: paddle.Tensor,
         gate: nn.Layer,
         topk_ids_hookfunc: Callable = None,
     ) -> paddle.Tensor:
         """
-        Paddle compute Fused MoE.
+        Apply the TP prefill method.
         """
-        # for all layer.ep_size and layer.tp_size
-        return self.apply_tp(layer, x, gate, topk_ids_hookfunc=topk_ids_hookfunc)
+        raise NotImplementedError
 
     def apply_ep_prefill(
         self,
@@ -96,7 +95,7 @@ def apply_ep_decode(
         """
         raise NotImplementedError
 
-    def apply_tp(
+    def apply(
         self,
         layer: nn.Layer,
         x: paddle.Tensor,
@@ -209,17 +208,18 @@ def _extract_descale_tensor(key_template, logical_expert_ids):
         down_proj_weight = paddle.stack(down_proj_weights, axis=0)
         up_gate_proj_weight_scale = paddle.stack(up_gate_proj_weight_scale, axis=0)
         down_proj_weight_scale = paddle.stack(down_proj_weight_scale, axis=0)
+        down_proj_in_scale = paddle.stack(down_proj_in_scale, axis=0)
 
         name_tensor_map = {
             "up_gate_proj_weight": up_gate_proj_weight,
             "down_proj_weight": down_proj_weight,
             "up_gate_proj_weight_scale": up_gate_proj_weight_scale,
             "down_proj_weight_scale": down_proj_weight_scale,
             "up_gate_proj_in_scale": up_gate_proj_in_scale,
+            "down_proj_in_scale": down_proj_in_scale,
         }
         for name, tensor in name_tensor_map.items():
             getattr(layer, name).set_value(tensor)
-        setattr(layer, "down_proj_in_scale", down_proj_in_scale)
 
     def create_weights(self, layer: nn.Layer, **extra_weight_attrs):
         """
@@ -266,6 +266,15 @@ def create_weights(self, layer: nn.Layer, **extra_weight_attrs):
                 default_initializer=paddle.nn.initializer.Constant(0),
             ),
         )
+        setattr(
+            layer,
+            "down_proj_in_scale",
+            layer.create_parameter(
+                shape=[layer.num_local_experts, 1],
+                dtype=self.default_dtype,
+                default_initializer=paddle.nn.initializer.Constant(0),
+            ),
+        )
 
         # weight_scales
         setattr(
@@ -321,17 +330,17 @@ def init_ep(self, layer: nn.Layer) -> None:
         """
         return
 
-    def apply(
+    def apply_tp(
         self,
         layer: nn.Layer,
         x: paddle.Tensor,
         gate: nn.Layer,
         topk_ids_hookfunc: Callable = None,
     ) -> paddle.Tensor:
         """
-        Paddle compute Fused MoE.
+        Apply the TP decoder method.
         """
-        return self.apply_tp(layer, x, gate, topk_ids_hookfunc=topk_ids_hookfunc)
+        raise NotImplementedError
 
     def apply_ep_prefill(
         self,
@@ -357,7 +366,7 @@ def apply_ep_decode(
         """
         raise NotImplementedError
 
-    def apply_tp(
+    def apply(
         self,
         layer: nn.Layer,
         x: paddle.Tensor,

fastdeploy/worker/hpu_model_runner.py

@@ -326,6 +326,24 @@ def fused_mlp_forward(
 from fastdeploy.model_executor.models.qwen3 import Qwen3Attention
 
 
+def process_scales_alignment_after_loading(layer, alignment=0x80) -> None:
+    aligned_attrs = ["down_proj_in_scale"]
+    for attr in aligned_attrs:
+        if hasattr(layer, attr):
+            value = getattr(layer, attr)
+            if value.dim() != 2 or value.shape[1] != 1:
+                raise ValueError(f"{attr} must be [num_experts, 1], but got {value.shape}")
+
+            dtype_size = value.element_size()
+            elements_per_slot = alignment // dtype_size
+            # special 3D shape [..., 1, align_with_padding] as a mark for HPU alignment
+            padded_tensor = paddle.zeros(shape=[value.shape[0], 1, elements_per_slot], dtype=value.dtype)
+            padded_tensor[:, 0, 0] = value[:, 0]
+
+            delattr(layer, attr)
+            setattr(layer, attr, padded_tensor)
+
+
 def convert_model(model, measurement_mode=False, init_done=False):
     """ """
     if measurement_mode:
@@ -354,6 +372,7 @@ def convert_model(model, measurement_mode=False, init_done=False):
                 module.forward = types.MethodType(fused_attention_forward, module)
             if isinstance(module, FusedMoE):
                 module.measurement_mode = measurement_mode
+                process_scales_alignment_after_loading(module)
 
     return model