add block scale moe python api for asm_moe (#341)

* ck fuse moe gemm stage1 with act

* ck fuse moe gemm stage1 with act

* fp8 switch to ck

* move to ck develop

* cutting some moe ut cases

* fix bug

* fix profile bugs

* revert changes of 'get_2stage_cfgs'

* re-trigger CI

* moe i8 switch ck

* Retrigger CI

* add block scale moe python api for asm_moe

* blockscale moe

* fix blockscale bugs

Author: lalala-sh

aiter/fused_moe_bf16_asm.py

@@ -54,6 +54,7 @@ def asm_moe(hidden_states,
             fc2_smooth_scale=None,  # [expert(local_expert:EP), 1, inter_dim]
             a16=False,
             per_tensor_quant_scale=None,
+            block_shape=None,
             expert_mask=None,
             activation = ActivationType.Silu
             ):
@@ -93,7 +94,41 @@ def asm_moe(hidden_states,
         else:
             raise ValueError(
                 f"Invalid args: {w1.dtype} {w1.shape=} {w2.shape=}")
-
+    elif block_shape is not None:
+        assert dtype == torch.bfloat16, "asm_moe for block_scale only support bfloat16 hidden_states"
+        assert block_shape == (
+            128, 128), "asm_moe for block_scale only support (128, 128)"
+        assert w1.dtype == torch.float8_e4m3fnuz, "asm_moe for block_scale only support float8_e4m3fnuz weight"
+        assert w2.shape[2] * 2 == w1.shape[1], "aiter moe for block_scale only support g1u1"
+        scale_blk_n, scale_blk_k = block_shape
+        hidden_states = hidden_states.view(M *
+                                           model_dim//scale_blk_k, scale_blk_k)
+
+        a1_q, a1_scale = pertoken_quant(
+            hidden_states.view(-1, model_dim // scale_blk_k, scale_blk_k), quant_dtype=torch.float8_e4m3fnuz
+        )
+        a1_q = a1_q.view(-1, model_dim)
+        a1_scale = a1_scale.squeeze(-1).t().contiguous()
+
+
+        scale_blk_n, scale_blk_k = block_shape
+        aiter.fmoe_fp8_blockscale_g1u1(
+            moe_buf,
+            a1_q,
+            w1,
+            w2,
+            sorted_ids,
+            sorted_weights,
+            sorted_expert_ids,
+            num_valid_ids,
+            topk,
+            a1_scale,
+            fc1_scale,
+            fc2_scale,
+            scale_blk_n,
+            scale_blk_k,
+            None,
+        )
     else:
         # a8w8 fmoe, opt: smooth quant
         a8_type = w1.dtype if w1.dtype != torch.int32 and w1.dtype != torch.uint32 else torch.float8_e4m3fnuz

op_tests/test_moe.py

@@ -90,6 +90,7 @@ def asm_moe_test(
         a16,
         None,
         None,
+        None,
         activation,
     )