Add xe2 block_fp8 grouped_gemm kernel

benchmark/benchmark_cutlass_fused_moe.py

@@ -45,7 +45,7 @@ def calculate_memory_usage(m, n, k, num_experts, x_dtype, w_dtype=None):
 
 
 def make_fused_moe_input(config):
-    mnk, e, topk, x_dtype, w_dtype, has_bias = config
+    mnk, e, topk, x_dtype, w_dtype, has_bias, is_block_fp8 = config
     m, n, k = mnk
     input_len = m
     hidden_size = k
@@ -83,7 +83,55 @@ def make_fused_moe_input(config):
     flat_expert_indices = expert_indices.view(-1)
     flat_expert_weights = expert_scores.view(-1, 1)
 
-    if w_dtype is not None:
+    if w_dtype is not None and is_block_fp8:
+        # 128x128 block-wise fp8 weights. Weights are laid out [E, N, K] here
+        # (natural orientation) and transposed to the kernel's [E, K, N] below,
+        # so the kernel scales are the transpose of the natural block grid
+        # [E, N // 128, K // 128] -> [E, K // 128, N // 128].
+        block_size = 128
+        n13 = 2 * intermediate_size
+        assert (hidden_size % block_size == 0
+                and intermediate_size % block_size == 0), \
+            "block fp8 requires hidden_size and intermediate_size " \
+            "divisible by 128"
+
+        w13_blk_scales = torch.pow(
+            2.0,
+            torch.randint(-3,
+                          4, (num_experts, n13 // block_size,
+                              hidden_size // block_size),
+                          device=DEVICE).float())
+        w2_blk_scales = torch.pow(
+            2.0,
+            torch.randint(-3,
+                          4, (num_experts, hidden_size // block_size,
+                              intermediate_size // block_size),
+                          device=DEVICE).float())
+
+        w13_fp8 = torch.empty_like(w13, dtype=w_dtype)
+        w2_fp8 = torch.empty_like(w2, dtype=w_dtype)
+        ref_w13 = torch.empty_like(w13, dtype=x_dtype)
+        ref_w2 = torch.empty_like(w2, dtype=x_dtype)
+        for i in range(num_experts):
+            # Quantize by dividing each 128x128 block by its scale before the
+            # fp8 cast (mirroring the per-tensor scaled_fp8_quant path), so the
+            # dequantized weight (fp8 * scale) reconstructs the original weight
+            # and output magnitudes stay in the same range as the bf16 / per-
+            # tensor reference.
+            s13 = w13_blk_scales[i].repeat_interleave(
+                block_size, dim=0).repeat_interleave(block_size, dim=1)
+            w13_fp8[i] = (w13[i] / s13).to(w_dtype)
+            ref_w13[i] = w13_fp8[i].to(x_dtype) * s13
+            s2 = w2_blk_scales[i].repeat_interleave(
+                block_size, dim=0).repeat_interleave(block_size, dim=1)
+            w2_fp8[i] = (w2[i] / s2).to(w_dtype)
+            ref_w2[i] = w2_fp8[i].to(x_dtype) * s2
+        w13 = w13_fp8
+        w2 = w2_fp8
+        # kernel scales: [E, K // 128, N // 128]
+        w13_scales = w13_blk_scales.transpose(1, 2).contiguous()
+        w2_scales = w2_blk_scales.transpose(1, 2).contiguous()
+    elif w_dtype is not None:
         w13_fp8 = torch.empty_like(w13, dtype=w_dtype)
         w2_fp8 = torch.empty_like(w2, dtype=w_dtype)
 
@@ -126,7 +174,7 @@ def make_fused_moe_input(config):
 
 
 def calculate_diff(config):
-    _, e, topk, x_dtype, w_dtype, _ = config
+    _, e, topk, x_dtype, w_dtype, _, is_block_fp8 = config
     ref_a, ref_w13, w13_bias, ref_w2, w2_bias, flat_expert_weights, \
         flat_expert_indices, a, w13, w13_scales, w2, w2_scales, \
             expert_scores, expert_indices = make_fused_moe_input(config)
@@ -147,7 +195,9 @@ def calculate_diff(config):
                            n_experts_per_token=topk,
                            activation="silu",
                            num_experts=e,
-                           is_fp8=(w_dtype is not None))
+                           is_fp8=(w_dtype is not None
+                                   and not is_block_fp8),
+                           is_block_fp8=is_block_fp8)
     if x_dtype == torch.float16:
         rtol = 1e-2
         atol = 1e-2
@@ -168,7 +218,7 @@ def get_benchmark(iterations):
         triton.testing.Benchmark(
             x_names=[
                 "m", "n", "k", "num_experts", "topk", "x_dtype", "w_dtype",
-                "has_bias"
+                "has_bias", "is_block_fp8"
             ],
             x_vals=[(*tuple(c)[0], *tuple(c)[1:]) for c in configs],
             line_arg="provider",
@@ -202,11 +252,12 @@ def benchmark(m,
                   x_dtype,
                   w_dtype,
                   has_bias,
+                  is_block_fp8,
                   provider,
                   iterations=iterations):
         print(f"Running config: {m, n, k, num_experts, topk, \
                                 x_dtype, w_dtype, \
-                                has_bias}, Provider: {provider}",
+                                has_bias, is_block_fp8}, Provider: {provider}",
               flush=True)
         total_latency = 0.0
         ms = 0.0
@@ -218,7 +269,7 @@ def benchmark(m,
             _, a, w13, w13_scales, w2, w2_scales, \
                 expert_scores, expert_indices = make_fused_moe_input(
                     config=((m, n, k), num_experts,
-                            topk, x_dtype, w_dtype, has_bias))
+                            topk, x_dtype, w_dtype, has_bias, is_block_fp8))
 
         if provider == "vllm":
             start_event = torch.xpu.Event(enable_timing=True)
@@ -236,7 +287,9 @@ def benchmark(m,
                               n_experts_per_token=topk,
                               activation="silu",
                               num_experts=num_experts,
-                              is_fp8=(w_dtype is not None))
+                              is_fp8=(w_dtype is not None
+                                      and not is_block_fp8),
+                              is_block_fp8=is_block_fp8)
             start_event.record()
             for index in range(5, iterations):
                 xpu_fused_moe(hidden_states=a,
@@ -251,7 +304,9 @@ def benchmark(m,
                               n_experts_per_token=topk,
                               activation="silu",
                               num_experts=num_experts,
-                              is_fp8=(w_dtype is not None))
+                              is_fp8=(w_dtype is not None
+                                      and not is_block_fp8),
+                              is_block_fp8=is_block_fp8)
             end_event.record()
             torch.xpu.synchronize()
             total_latency = start_event.elapsed_time(end_event)
@@ -297,7 +352,8 @@ def benchmark(m,
                     n_experts_per_token=topk,
                     activation="silu",
                     num_experts=num_experts,
-                    is_fp8=(w_dtype is not None),
+                    is_fp8=(w_dtype is not None and not is_block_fp8),
+                    is_block_fp8=is_block_fp8,
                     start_event_remap=remap_se[i] if i is not None else None,
                     end_event_remap=remap_ee[i] if i is not None else None,
                     start_event_gemm1=gemm1_se[i] if i is not None else None,

benchmark/src/fused_moe_interface_.py

@@ -28,6 +28,7 @@ def xpu_fused_moe_CalKernelTime(hidden_states,
                                 is_fp8=False,
                                 is_int4=False,
                                 is_mxfp4=False,
+                                is_block_fp8=False,
                                 start_event_remap=None,
                                 end_event_remap=None,
                                 start_event_gemm1=None,
@@ -107,7 +108,7 @@ def xpu_fused_moe_CalKernelTime(hidden_states,
                                dtype=hidden_states.dtype,
                                device=hidden_states.device)
 
-    if not is_fp8 and not is_int4 and not is_mxfp4:
+    if not is_fp8 and not is_int4 and not is_mxfp4 and not is_block_fp8:
         gemm1_scales = None
         gemm2_scales = None
     else:
@@ -172,7 +173,8 @@ def xpu_fused_moe_CalKernelTime(hidden_states,
         K=hidden_size,
         num_experts=num_experts,
         is_B_int4=is_int4,
-        is_B_mxfp4=is_mxfp4)
+        is_B_mxfp4=is_mxfp4,
+        is_B_fp8block=is_block_fp8)
     if end_event_gemm1 is not None:
         end_event_gemm1.record()
     active_experts1 = (rows_per_expert > 0).sum().item()
@@ -212,7 +214,8 @@ def xpu_fused_moe_CalKernelTime(hidden_states,
         K=inter_size,
         num_experts=num_experts,
         is_B_int4=is_int4,
-        is_B_mxfp4=is_mxfp4)
+        is_B_mxfp4=is_mxfp4,
+        is_B_fp8block=is_block_fp8)
 
     if end_event_gemm2 is not None:
         end_event_gemm2.record()

benchmark/src/get_model_config.py

@@ -27,9 +27,16 @@ def gen_cutlass_fused_moe_correctness_configs():
     x_dtype = [torch.float16, torch.bfloat16]
     w_dtype = [torch.float8_e5m2, torch.float8_e4m3fn, None]
     has_bias = [True, False]
-
-    configs = list(
-        itertools.product(mnk, experts, topk, x_dtype, w_dtype, has_bias))
+    # block_fp8 = True uses 128x128 block-wise fp8 weight scales instead of a
+    # single per-expert (per-tensor) scale. It is only valid for fp8 weights
+    # whose N and K are both divisible by 128.
+    is_block = [False, True]
+
+    configs = [
+        (m, e, t, x, w, b, blk) for m, e, t, x, w, b, blk in itertools.product(
+            mnk, experts, topk, x_dtype, w_dtype, has_bias, is_block)
+        if not (blk and (w is None or m[1] % 128 != 0 or m[2] % 128 != 0))
+    ]
     return configs
 
 
@@ -38,6 +45,9 @@ def gen_cutlass_fused_moe_perf_configs():
     x_dtype = [torch.float16, torch.bfloat16]
     w_dtype = [torch.float8_e5m2, torch.float8_e4m3fn, None]
     has_bias = [True, False]
+    # block_fp8 = True uses 128x128 block-wise fp8 weight scales; only valid for
+    # fp8 weights with N and K both divisible by 128.
+    is_block = [False, True]
     input_lens = [1, 4, 16, 1024, 8192]
 
     for model in model_lists:
@@ -53,9 +63,13 @@ def gen_cutlass_fused_moe_perf_configs():
 
         moe_top_k = model_config["moe_config"]["moe_top_k"]
         num_experts = model_config["num_groups"]
-        configs += list(
-            itertools.product(mnk, [num_experts],
-                              [moe_top_k], x_dtype, w_dtype, has_bias))
+        configs += [
+            (mk, ne, tk, x, w, b, blk)
+            for mk, ne, tk, x, w, b, blk in itertools.product(
+                mnk, [num_experts], [moe_top_k], x_dtype, w_dtype, has_bias,
+                is_block)
+            if not (blk and (w is None or mk[1] % 128 != 0 or mk[2] % 128 != 0))
+        ]
 
     # Hardcoded model shapes (n, k, num_experts, topk) for various TP sizes.
     # Config tuple (m, n, k) produces:
@@ -99,16 +113,21 @@ def gen_cutlass_fused_moe_perf_configs():
     for n, k, num_experts, topk in hardcoded_model_shapes:
         mnk = list(zip(input_lens, [n] * len(input_lens),
                        [k] * len(input_lens)))
-        configs += list(
-            itertools.product(mnk, [num_experts],
-                              [topk], x_dtype, w_dtype, has_bias))
+        configs += [
+            (mk, ne, tk, x, w, b, blk)
+            for mk, ne, tk, x, w, b, blk in itertools.product(
+                mnk, [num_experts], [topk], x_dtype, w_dtype, has_bias,
+                is_block)
+            if not (blk and (w is None or mk[1] % 128 != 0 or mk[2] % 128 != 0))
+        ]
 
     configs = set(configs)  # remove duplicates
 
     def sort_key(x):
-        (m, n, k), moe_topk, topk_, x_dtype_, w_dtype_, bias_ = x
+        (m, n, k), moe_topk, topk_, x_dtype_, w_dtype_, bias_, blk_ = x
 
-        return (m, n, k, moe_topk, topk_, str(x_dtype_), str(w_dtype_), bias_)
+        return (m, n, k, moe_topk, topk_, str(x_dtype_), str(w_dtype_), bias_,
+                blk_)
 
     configs = sorted(configs, key=sort_key)
     return configs

csrc/xpu/grouped_gemm/grouped_gemm_interface.cpp

@@ -20,8 +20,16 @@ torch::Tensor cutlass_grouped_gemm_interface(
     int64_t N,
     int64_t K,
     int64_t num_experts) {
+  const auto B_dtype = ptr_B.dtype();
+  const bool is_B_fp8block =
+      (B_dtype == at::kFloat8_e4m3fn || B_dtype == at::kFloat8_e5m2) &&
+      ptr_B_scale.has_value() && ptr_B_scale->dim() == 3;
   if (vllm::xpu::force_xe_default_kernel()) {
 #ifdef VLLM_XPU_ENABLE_XE_DEFAULT
+    TORCH_CHECK(
+        !is_B_fp8block,
+        "Block-wise FP8 grouped gemm is not supported by the XE default "
+        "kernel.");
     int64_t groups = num_experts;
     return cutlass_grouped_gemm_xe_default(
         ptr_A, ptr_B, ptr_bias, ptr_D, rows_per_expert, N, K, groups);
@@ -49,6 +57,10 @@ torch::Tensor cutlass_grouped_gemm_interface(
 #endif
   } else {
 #ifdef VLLM_XPU_ENABLE_XE_DEFAULT
+    TORCH_CHECK(
+        !is_B_fp8block,
+        "Block-wise FP8 grouped gemm is not supported by the XE default "
+        "kernel.");
     int64_t groups = num_experts;
     return cutlass_grouped_gemm_xe_default(
         ptr_A, ptr_B, ptr_bias, ptr_D, rows_per_expert, N, K, groups);