redirect asm_moe_tkw1 call to fused_moe in order to force kernel tuning

aiter/configs/tuned_fmoe.csv

@@ -772,3 +772,4 @@ cu_num,token,model_dim,inter_dim,expert,topk,act_type,dtype,q_dtype_a,q_dtype_w,
 80,256,7168,256,257,9,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fnuz,torch.float8_e4m3fnuz,QuantType.per_Token,1,0,32,0,536.7655,_ZN5aiter45fmoe_bf16_pertokenFp8_g1u1_vs_silu_1tg_32x256E,0.0%,0.0,Null,0.0%,536.7655,1,47.26,2646.03
 80,512,7168,256,257,9,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fnuz,torch.float8_e4m3fnuz,QuantType.per_Token,1,0,32,0,560.4425,_ZN5aiter45fmoe_bf16_pertokenFp8_g1u1_vs_silu_1tg_32x256E,0.0%,0.0,Null,0.0%,560.4425,1,90.53,2544.06
 80,1024,7168,256,257,9,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fnuz,torch.float8_e4m3fnuz,QuantType.per_Token,1,0,32,0,827.4898,_ZN5aiter45fmoe_bf16_pertokenFp8_g1u1_vs_silu_1tg_32x256E,0.0%,0.0,Null,0.0%,827.4898,1,122.62,1736.35
+256,16,5120,1024,128,1,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float8_e4m3fn,QuantType.per_Token,1,1,32,0,0.0,_ZN5aiter43fmoe_bf16_pertokenInt8_g1u1_tkw1_silu_32x64E,0.0%,0.0,Null,0,0.0,1,0.0,0.0

aiter/fused_moe.py

@@ -104,6 +104,8 @@ def fused_moe(
     num_local_tokens: Optional[torch.tensor] = None,
     moe_sorting_dispatch_policy=0,
     dtype=None,
+    a16=False,
+    per_tensor_quant_scale=None,
     # following for cktile support
     hidden_pad=0,
     intermediate_pad=0,
@@ -130,6 +132,8 @@ def fused_moe(
         num_local_tokens=num_local_tokens,
         moe_sorting_dispatch_policy=moe_sorting_dispatch_policy,
         dtype=dtype,
+        a16=a16,
+        per_tensor_quant_scale=per_tensor_quant_scale,
         hidden_pad=hidden_pad,
         intermediate_pad=intermediate_pad,
         bias1=bias1,
@@ -157,6 +161,8 @@ def fused_moe_fake(
     num_local_tokens: Optional[torch.Tensor] = None,
     moe_sorting_dispatch_policy: bool = 0,
     dtype: Optional[torch.dtype] = None,
+    a16: bool = False,
+    per_tensor_quant_scale: torch.Tensor = None,
 ) -> torch.Tensor:
     device = topk_ids.device
     M, topk = topk_ids.shape
@@ -187,6 +193,8 @@ def fused_moe_(
     num_local_tokens: Optional[torch.Tensor] = None,
     moe_sorting_dispatch_policy: bool = 0,
     dtype: Optional[torch.dtype] = None,
+    a16: bool = False,
+    per_tensor_quant_scale: torch.Tensor = None,
     hidden_pad: int = 0,
     intermediate_pad: int = 0,
     bias1: Optional[torch.Tensor] = None,
@@ -253,10 +261,7 @@ def fused_moe_(
         moe_sorting_dispatch_policy,
     )
 
-    if metadata.run_1stage:
-        assert (
-            doweight_stage1 == False
-        ), "doweight_stage1 not support in fused_moe_1stage"
+    if metadata.run_1stage and not doweight_stage1:
         return metadata.stage1(
             hidden_states,
             w1,
@@ -279,6 +284,26 @@ def fused_moe_(
             a2_scale=a2_scale,
             num_local_tokens=num_local_tokens,
         )
+    elif metadata.run_1stage and doweight_stage1:
+        return metadata.stage1(
+            hidden_states,
+            w1,
+            w2,
+            topk_ids,
+            sorted_ids,
+            sorted_weights,
+            sorted_expert_ids,
+            num_valid_ids,
+            moe_buf,
+            w1_scale,
+            w2_scale,
+            a1_scale,
+            a2_scale,
+            a16,
+            per_tensor_quant_scale,
+            expert_mask,
+            activation,
+        )
     else:
         return fused_moe_2stages(
             hidden_states,
@@ -429,6 +454,156 @@ def fused_moe_1stage(
     return moe_buf
 
 
+def fused_moe_stage1_tkw1(
+    hidden_states,
+    w1,  # [expert(local_expert:EP), inter_dim*2, dim] N,K
+    w2,  # [expert(local_expert:EP), dim, inter_dim]
+    topk_ids,
+    sorted_ids,
+    sorted_weights,
+    sorted_expert_ids,
+    num_valid_ids,
+    moe_buf,
+    # following for int8 quant
+    fc1_scale=None,  # [expert(local_expert:EP), inter_dim, 1]
+    fc2_scale=None,  # [expert(local_expert:EP), model_dim, 1]
+    fc1_smooth_scale=None,  # [expert(local_expert:EP), 1, model_dim]
+    fc2_smooth_scale=None,  # [expert(local_expert:EP), 1, inter_dim]
+    a16=False,
+    per_tensor_quant_scale=None,
+    expert_mask=None,
+    activation=ActivationType.Silu,
+    kernelName: str = "",
+):
+    E, model_dim, inter_dim = w2.shape
+    M, topk = topk_ids.shape
+    device = topk_ids.device
+    lastdim_mul = 8 if w1.dtype in {dtypes.i32, torch.uint32} else 1
+
+    if fc1_scale is None:
+        # pure bf16
+        aiter.fmoe(
+            moe_buf,
+            hidden_states,
+            w1,
+            w2,
+            sorted_ids,
+            sorted_weights,
+            sorted_expert_ids,
+            num_valid_ids,
+            topk,
+        )
+    elif a16:
+        # a16w8 smooth quant fmoe
+        if w1.dtype == dtypes.fp8 and inter_dim * 2 == w1.shape[1]:
+            aiter.fmoe_fp8_g1u1_a16(
+                moe_buf,
+                hidden_states,
+                w1,
+                w2,
+                sorted_ids,
+                sorted_weights,
+                sorted_expert_ids,
+                num_valid_ids,
+                topk,
+                fc1_scale,
+                fc2_scale,
+                fc1_smooth_scale,
+                fc2_smooth_scale,
+            )
+        elif w1.dtype == dtypes.i8 and inter_dim == w1.shape[1]:
+            aiter.fmoe_int8_g1u0_a16(
+                moe_buf,
+                hidden_states,
+                w1,
+                w2,
+                sorted_ids,
+                sorted_weights,
+                sorted_expert_ids,
+                num_valid_ids,
+                topk,
+                fc1_scale,
+                fc2_scale,
+                fc1_smooth_scale,
+                fc2_smooth_scale,
+            )
+        else:
+            raise ValueError(f"Invalid args: {w1.dtype} {w1.shape=} {w2.shape=}")
+
+    else:
+        # a8w8 fmoe, opt: smooth quant
+        a8_type = (
+            w1.dtype
+            if w1.dtype != dtypes.i32 and w1.dtype != torch.uint32
+            else dtypes.fp8
+        )
+        if fc1_smooth_scale is not None:
+            a8 = torch.empty((topk * M, model_dim), dtype=a8_type, device=device)
+            a8_scale = torch.empty((topk * M), dtype=dtypes.fp32, device=device)
+
+            # moe_smoothquant_fwd need topk_ids which contains local_expert_id
+            if expert_mask is not None:
+                local_expert_hash = expert_mask.cumsum(0, dtype=dtypes.i32)
+                local_expert_hash[local_expert_hash > 0] -= 1
+                topk_ids = local_expert_hash[topk_ids]
+
+            aiter.moe_smoothquant_fwd(
+                a8, hidden_states, fc1_smooth_scale, topk_ids, a8_scale
+            )
+        else:
+            if (
+                w1.dtype == dtypes.fp8
+                or w1.dtype == dtypes.i32
+                and w1.dtype == torch.uint32
+            ):
+                a8 = torch.empty((M, model_dim), dtype=a8_type, device=device)
+                a8_scale = torch.empty(M, dtype=dtypes.fp32, device=device)
+                if per_tensor_quant_scale is None:
+                    aiter.dynamic_per_token_scaled_quant(a8, hidden_states, a8_scale)
+                else:
+                    aiter.static_per_tensor_quant(
+                        a8, hidden_states, per_tensor_quant_scale
+                    )
+                    a8_scale.fill_(per_tensor_quant_scale)
+            elif w1.dtype == dtypes.i8:
+                a8 = torch.empty((M, model_dim), dtype=w1.dtype, device=device)
+                a8_scale = torch.empty(M, dtype=dtypes.fp32, device=device)
+                fc1_smooth_scale = torch.ones(
+                    model_dim, dtype=dtypes.fp32, device=device
+                )
+                aiter.smoothquant_fwd(a8, hidden_states, fc1_smooth_scale, a8_scale)
+            else:
+                logger.warning("FMOE fall into pure torch quant...")
+                a8, a8_scale = aiter.pertoken_quant(hidden_states, quant_dtype=w1.dtype)
+        if w2.shape[2] * 2 * lastdim_mul == w1.shape[1]:
+            fmoe_func = aiter.fmoe_g1u1_tkw1
+
+        else:
+            raise ValueError(
+                f"Invalid MoE weight: {w1.shape=} {w2.shape=} {lastdim_mul}"
+            )
+
+        fmoe_func(
+            moe_buf,
+            a8,
+            w1,
+            w2,
+            sorted_ids,
+            sorted_weights,
+            sorted_expert_ids,
+            num_valid_ids,
+            topk,
+            a8_scale,
+            fc1_scale,
+            fc2_scale,
+            kernelName,
+            fc2_smooth_scale,
+            activation,
+        )
+        #   fc2_smooth_scale)
+    return moe_buf
+
+
 @functools.lru_cache(maxsize=1024)
 def get_block_size_M(token, topk, expert, inter_dim):
     cu_num = get_cu_num()
@@ -468,6 +643,9 @@ def get_block_size_M(token, topk, expert, inter_dim):
     {
         (ActivationType.Silu,    QuantType.per_1x32,   dtypes.bf16,   dtypes.fp4x2,  dtypes.fp4x2,    True) : aiter.fmoe_g1u1,
         (ActivationType.Silu,   QuantType.per_1x128,   dtypes.bf16,     dtypes.fp8,    dtypes.fp8,    True) : aiter.fmoe_fp8_blockscale_g1u1,
+        (ActivationType.Silu,   QuantType.per_Token,   dtypes.bf16,    dtypes.bf16,   dtypes.bf16,   False) : aiter.fmoe,
+        (ActivationType.Silu,   QuantType.per_Token,   dtypes.bf16,     dtypes.i8,     dtypes.i8,    False) : aiter.fmoe_int8_g1u0_a16,
+        (ActivationType.Silu,   QuantType.per_Token,   dtypes.bf16,     dtypes.fp8,    dtypes.fp8,    True) : aiter.fmoe_g1u1_tkw1,
     }
 }
 # fmt: on
@@ -618,6 +796,19 @@ def FinalFunc():
                 run_1stage = token > 32
             elif q_type != QuantType.per_1x32:
                 run_1stage = token < 256
+        elif (
+            doweight_stage1
+            and (
+                activation,
+                q_type,
+                dtype,
+                q_dtype_a,
+                q_dtype_w,
+                use_g1u1,
+            )
+            in fused_moe_1stage_dict[get_gfx()]
+        ):
+            run_1stage = True
 
         block_m = (
             BLOCK_SIZE_M
@@ -639,7 +830,8 @@ def FinalFunc():
     logger.info(
         f"[fused_moe] using {'1stage' if run_1stage else '2stage'} {'default' if cfg is None else tag} for {keys} "
     )
-    if run_1stage:
+    if run_1stage and not doweight_stage1:
+        logger.info(f"[get_2stage_cfgs] run_1stage")
         return MOEMetadata(
             functools.partial(
                 fused_moe_1stage,
@@ -652,6 +844,15 @@ def FinalFunc():
             ksplit,
             run_1stage,
         )
+    elif run_1stage and doweight_stage1:
+        logger.info(f"[get_2stage_cfgs] run_1stage and doweight_stage1")
+        return MOEMetadata(
+            functools.partial(fused_moe_stage1_tkw1, kernelName=kernelName1),
+            None,
+            block_m,
+            ksplit,
+            run_1stage,
+        )
     if (
         dtype in [dtypes.bf16, dtypes.fp16]
         and q_type == QuantType.per_1x32