[QC] fused_moe blockwise fp8 w8a8 optimization (flagos-ai#2085)

Author: monellz

benchmark/test_vllm_perf.py

@@ -32,6 +32,7 @@ def is_cuda_available():
 
 
 CUDA_AVAILABLE = is_cuda_available()
+DEFAULT_BLOCK_SHAPE = [128, 128]
 
 
 def to_int8(tensor: torch.Tensor):
@@ -455,6 +456,117 @@ def _fp8_input_fn(self, config, dtype):
         )
 
 
+class FusedMoEFP8BlockwiseBenchmark(Benchmark):
+    """
+    Benchmark for fused_experts_impl with FP8 W8A8 block-wise quantization.
+
+    Weights are stored in FP8 E4M3 and accompanied by block scales.
+    Activations are dynamically quantized per-token per-group inside the kernel.
+    """
+
+    def __init__(self, op_name, torch_op, dtypes):
+        super().__init__(op_name=op_name, torch_op=torch_op, dtypes=dtypes)
+        self.block_shape = DEFAULT_BLOCK_SHAPE
+
+    def set_shapes(self, shape_file_path=None):
+        # (num_tokens, num_experts, hidden_size, intermediate_size, topk)
+        self.shapes = [
+            # Mixtral-like shapes
+            (1, 8, 4096, 14336, 2),
+            (4, 8, 4096, 14336, 2),
+            (16, 8, 4096, 14336, 2),
+            (64, 8, 4096, 14336, 2),
+            (128, 8, 4096, 14336, 2),
+            (256, 8, 4096, 14336, 2),
+            (512, 8, 4096, 14336, 2),
+            # DeepSeek-V3-like shapes (TP=8 shard)
+            (1, 256, 7168, 2048, 8),
+            (4, 256, 7168, 2048, 8),
+            (16, 256, 7168, 2048, 8),
+            (64, 256, 7168, 2048, 8),
+            (128, 256, 7168, 2048, 8),
+            (256, 256, 7168, 2048, 8),
+            # Qwen3.5-397B-A17B
+            (1, 512, 4096, 1024, 10),
+            (4, 512, 4096, 1024, 10),
+            (16, 512, 4096, 1024, 10),
+            (64, 512, 4096, 1024, 10),
+            (128, 512, 4096, 1024, 10),
+            (256, 512, 4096, 1024, 10),
+        ]
+
+    def get_input_iter(self, cur_dtype):
+        del cur_dtype
+        for config in self.shapes:
+            yield from self._fp8_blockwise_input_fn(config)
+
+    def _fp8_blockwise_input_fn(self, config):
+        num_tokens, num_experts, hidden_size, intermediate_size, topk = config
+        block_n, block_k = self.block_shape
+        device = flag_gems.device
+        dtype = torch.bfloat16
+
+        hidden_states = torch.randn(num_tokens, hidden_size, device=device, dtype=dtype)
+        w1_fp8 = (
+            torch.randn(
+                num_experts,
+                intermediate_size * 2,
+                hidden_size,
+                device=device,
+                dtype=torch.bfloat16,
+            )
+            * (1.0 / hidden_size**0.5)
+        ).to(torch.float8_e4m3fn)
+        w2_fp8 = (
+            torch.randn(
+                num_experts,
+                hidden_size,
+                intermediate_size,
+                device=device,
+                dtype=torch.bfloat16,
+            )
+            * (1.0 / intermediate_size**0.5)
+        ).to(torch.float8_e4m3fn)
+
+        w1_scale = (
+            torch.rand(
+                num_experts,
+                ceil(intermediate_size * 2 / block_n),
+                ceil(hidden_size / block_k),
+                device=device,
+                dtype=torch.float32,
+            )
+            + 0.01
+        )
+        w2_scale = (
+            torch.rand(
+                num_experts,
+                ceil(hidden_size / block_n),
+                ceil(intermediate_size / block_k),
+                device=device,
+                dtype=torch.float32,
+            )
+            + 0.01
+        )
+
+        gating = torch.randn(
+            num_tokens, num_experts, device=device, dtype=torch.float32
+        )
+        topk_weights, topk_ids = torch.topk(torch.softmax(gating, dim=-1), topk, dim=-1)
+        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
+        topk_weights = topk_weights.to(torch.float32)
+
+        yield (
+            hidden_states,
+            w1_fp8,
+            w2_fp8,
+            w1_scale,
+            w2_scale,
+            topk_weights,
+            topk_ids,
+        )
+
+
 def _vllm_fused_moe_fp8_wrapper(
     hidden_states, w1, w2, topk_weights, topk_ids, w1_scale, w2_scale
 ):
@@ -489,6 +601,42 @@ def _gems_fused_moe_fp8_wrapper(
     )
 
 
+def _vllm_fused_moe_fp8_blockwise_wrapper(
+    hidden_states, w1, w2, w1_scale, w2_scale, topk_weights, topk_ids
+):
+    """Wrapper to call vllm fused_experts_impl with block-wise FP8."""
+    return vllm_fused_experts_impl(
+        hidden_states.clone(),
+        w1,
+        w2,
+        topk_weights,
+        topk_ids,
+        inplace=False,
+        activation="silu",
+        use_fp8_w8a8=True,
+        w1_scale=w1_scale,
+        w2_scale=w2_scale,
+        block_shape=DEFAULT_BLOCK_SHAPE,
+    )
+
+
+def _gems_fused_moe_fp8_blockwise_wrapper(
+    hidden_states, w1, w2, w1_scale, w2_scale, topk_weights, topk_ids
+):
+    """Wrapper to call FlagGems fused_experts_impl with block-wise FP8."""
+    return flag_gems.fused_experts_impl(
+        hidden_states,
+        w1,
+        w2,
+        topk_weights,
+        topk_ids,
+        use_fp8_w8a8=True,
+        w1_scale=w1_scale,
+        w2_scale=w2_scale,
+        block_shape=DEFAULT_BLOCK_SHAPE,
+    )
+
+
 @pytest.mark.fused_moe
 @pytest.mark.skipif(
     not (HAS_VLLM_FUSED_MOE and CUDA_AVAILABLE),
@@ -507,6 +655,24 @@ def test_perf_fused_moe_fp8_gems_vs_vllm():
     bench.run()
 
 
+@pytest.mark.fused_moe
+@pytest.mark.skipif(
+    not (HAS_VLLM_FUSED_MOE and CUDA_AVAILABLE),
+    reason="requires vLLM and NVIDIA Hopper architecture for FP8 blockwise",
+)
+def test_perf_fused_moe_fp8_blockwise_gems_vs_vllm():
+    """
+    Benchmark FlagGems vs vLLM fused_experts_impl with FP8 W8A8 block-wise quantization.
+    """
+    bench = FusedMoEFP8BlockwiseBenchmark(
+        op_name="fused_moe_fp8_blockwise_gems_vs_vllm",
+        torch_op=_vllm_fused_moe_fp8_blockwise_wrapper,
+        dtypes=[torch.bfloat16],
+    )
+    bench.set_gems(_gems_fused_moe_fp8_blockwise_wrapper)
+    bench.run()
+
+
 class FusedMoEINT8Benchmark(Benchmark):
     """
     Benchmark for fused_experts_impl with INT8 W8A8 quantization.

src/flag_gems/fused/fused_moe.py

@@ -188,9 +188,42 @@ def try_get_optimal_moe_config(
     top_k: int,
     dtype: str | None,
     M: int,
+    E: int,
     block_shape: list[int] | None = None,
 ) -> dict[str, int]:
     override_config: Optional[dict[str, Any]] = None
+
+    is_hopper = torch.cuda.is_available() and torch.cuda.get_device_capability() == (
+        9,
+        0,
+    )
+    if (
+        is_hopper
+        and dtype == "fp8_w8a8"
+        and block_shape is not None
+        and len(block_shape) == 2
+    ):
+        # Use heuristic config like hpc-ops
+        avg_tokens_per_expert = M * top_k // E
+        if avg_tokens_per_expert <= 16:
+            block_size_m = 16
+        elif avg_tokens_per_expert <= 32:
+            block_size_m = 32
+        elif avg_tokens_per_expert <= 48:
+            block_size_m = 48
+        else:
+            block_size_m = 64
+        config = {
+            "BLOCK_SIZE_M": block_size_m,
+            "BLOCK_SIZE_N": block_shape[0],
+            "BLOCK_SIZE_K": block_shape[1],
+            "GROUP_SIZE_M": 1,
+            "num_warps": 4,
+            "num_stages": 3,
+            "SWAP_AB": True,
+        }
+        override_config = config
+
     if override_config:
         config = override_config
     else:
@@ -508,6 +541,19 @@ def _fp8_quantize(
         assert len(block_shape) == 2
         block_k = block_shape[1]
         assert A.size(-1) % block_k == 0
+        if A.ndim == 2 and A.stride(-1) == 1:
+            from flag_gems.ops.per_token_group_quant_fp8 import (
+                per_token_group_quant_fp8,
+            )
+
+            return per_token_group_quant_fp8(
+                A,
+                group_size=block_k,
+                eps=eps,
+                dtype=fp8_dtype,
+                column_major_scales=False,
+                scale_ue8m0=False,
+            )
         orig_shape = A.shape
         A_flat = A.reshape(-1, A.size(-1))
         M, K = A_flat.shape
@@ -913,6 +959,7 @@ def fused_moe_kernel(
     use_int8_w8a16: tl.constexpr,
     per_channel_quant: tl.constexpr,
     HAS_BIAS: tl.constexpr,
+    SWAP_AB: tl.constexpr,
 ):
     """Fused MoE kernel: token × expert GEMM with quantization support."""
     # Map pid to C block (grouped ordering for L2 reuse)
@@ -1000,6 +1047,8 @@ def fused_moe_kernel(
         bias = tl.load(bias_ptrs, mask=(offs_bn < N), other=0.0)
     # Accumulate C block in fp32
     accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+    if SWAP_AB:
+        accumulator_nm = tl.zeros((BLOCK_SIZE_N, BLOCK_SIZE_M), dtype=tl.float32)
     for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
         a = tl.load(
             a_ptrs,
@@ -1016,9 +1065,16 @@ def fused_moe_kernel(
                 a_scale = tl.load(
                     a_scale_ptrs + offs_ks * stride_ask, mask=token_mask, other=0.0
                 )
-                b_scale = tl.load(b_scale_ptrs + offs_ks * stride_bsk)
-
-                accumulator += tl.dot(a, b) * a_scale[:, None] * b_scale[None, :]
+                if SWAP_AB:
+                    b_scale = tl.load(b_scale_ptrs + offs_ks * stride_bsk)
+                    accumulator_nm += (
+                        tl.dot(tl.trans(b), tl.trans(a))
+                        * b_scale[:, None]
+                        * a_scale[None, :]
+                    )
+                else:
+                    b_scale = tl.load(b_scale_ptrs + offs_ks * stride_bsk)
+                    accumulator += tl.dot(a, b) * a_scale[:, None] * b_scale[None, :]
             else:
                 if use_fp8_w8a8:
                     accumulator = tl.dot(a, b, acc=accumulator)
@@ -1029,6 +1085,9 @@ def fused_moe_kernel(
         a_ptrs += BLOCK_SIZE_K * stride_ak
         b_ptrs += BLOCK_SIZE_K * stride_bk
 
+    if SWAP_AB:
+        accumulator = tl.trans(accumulator_nm)
+
     # Dequantization
     if use_int8_w8a16:
         accumulator = accumulator * b_scale
@@ -1210,6 +1269,7 @@ def invoke_fused_moe_triton_kernel(
     if block_shape is not None:
         BLOCK_SIZE_K = min(BLOCK_SIZE_K, min(block_shape[0], block_shape[1]))
 
+    swap_AB = config.pop("SWAP_AB", False)
     fused_moe_kernel[grid](
         A,
         B,
@@ -1251,6 +1311,7 @@ def invoke_fused_moe_triton_kernel(
         naive_block_assignment=(sorted_token_ids is None),
         HAS_BIAS=HAS_BIAS,
         BLOCK_SIZE_K=BLOCK_SIZE_K,
+        SWAP_AB=swap_AB,
         **config,
     )
 
@@ -1431,6 +1492,7 @@ def fused_experts_impl(
         top_k_num,
         config_dtype,
         block_shape=block_shape,
+        E=E,
     )
 
     config = get_config_func(M)