fix: w4afp8 accuracy problem and rebase (#8752)

Signed-off-by: yangsijia.614 <yangsijia.614@bytedance.com>
Co-authored-by: Jinwu <ayrnb@users.noreply.github.com>

python/sglang/srt/layers/quantization/w4afp8.py

@@ -116,6 +116,8 @@ class W4AFp8MoEMethod(FusedMoEMethodBase):
         params_dtype: torch.dtype,
         **extra_weight_attrs,
     ):
+        from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported
+
         assert "weight_loader" in extra_weight_attrs
 
         # Fused gate_up_proj (column parallel)
@@ -144,6 +146,9 @@ class W4AFp8MoEMethod(FusedMoEMethodBase):
         layer.register_parameter("w2_weight", w2_weight)
         set_weight_attrs(w2_weight, extra_weight_attrs)
 
+        extra_weight_attrs.update(
+            {"quant_method": FusedMoeWeightScaleSupported.GROUP.value}
+        )
         w13_weight_scale = torch.nn.Parameter(
             torch.zeros(
                 num_experts,
@@ -274,8 +279,11 @@ class W4AFp8MoEMethod(FusedMoEMethodBase):
     def apply(
         self,
         layer: EPMoE,
-        hidden_states: torch.Tensor,
+        x: torch.Tensor,
         topk_output: TopKOutput,
+        activation: str = "silu",
+        apply_router_weight_on_input: bool = False,
+        routed_scaling_factor: Optional[float] = None,
         **kwargs,
     ) -> torch.Tensor:
 
@@ -284,19 +292,17 @@ class W4AFp8MoEMethod(FusedMoEMethodBase):
 
         topk_weights, topk_ids, _ = topk_output
         local_topk_ids = topk_ids
-        if layer.expert_map is not None:
-            "Translate info from expert_map to topk_ids"
-            local_topk_ids = torch.where(
-                layer.expert_map[topk_ids] != layer.num_experts,
-                layer.expert_map[topk_ids],
-                layer.num_experts,
-            )
+        local_topk_ids = torch.where(
+            topk_ids == -1,
+            layer.num_experts,
+            topk_ids,
+        )
 
-        return cutlass_w4a8_moe(
+        output = cutlass_w4a8_moe(
             layer.start_expert_id,
             layer.end_expert_id,
             layer.num_experts,
-            hidden_states,
+            x,
             layer.w13_weight,
             layer.w2_weight,
             layer.w13_weight_scale_inv,
@@ -318,3 +324,6 @@ class W4AFp8MoEMethod(FusedMoEMethodBase):
             layer.w13_input_scale,
             layer.w2_input_scale,
         )
+        if routed_scaling_factor is not None:
+            output *= routed_scaling_factor
+        return output