[CPU] Optimize routed FP8/MXFP4 MoE GEMM dispatch (#50949)
Signed-off-by: Li, Tianmu <tianmu.li@intel.com>
Tianmu Li committed
Aug 7, 2026 at 21:56 UTC
700d39b55813fce6fe5339b334beb55d2a39f598
3 files changed
+110
-5
csrc/cpu/sgl-kernels/gemm.h
+5
@@ -59,6 +59,11 @@ inline bool can_use_brgemm<at::Float8_e4m3fn>(int M) {
59
return brgemm_supported() && M > 4;
60
}
61
62
+template <typename T>
63
+inline bool can_use_brgemm(int M, int N) {
64
+ return brgemm_supported() && (can_use_brgemm<T>(M) || N != block_size_n());
65
+}
66
+
67
// work around compiler internal error
68
#define BLOCK_K 128 // 4 * TILE_K
69
csrc/cpu/sgl-kernels/moe_fp8.cpp
+14
-5
@@ -60,9 +60,6 @@ void fused_experts_fp_kernel_impl(
60
const int64_t stride_e = 2 * N * packed_K;
61
const int64_t stride_n = packed_K;
62
63
- int64_t avg_M = std::max(int64_t(1), M * topk / E);
64
- const bool use_brgemm = can_use_brgemm<packed_t>(avg_M);
65
-
63
int64_t B_tmp_size_per_thread = MAX_CACHE_BLOCK_SIZE * BLOCK_N * std::max(K, N);
64
65
// here we only parallel on half of 2N to fuse silu_and_mul with gemm
@@ -70,6 +67,7 @@ void fused_experts_fp_kernel_impl(
67
// get local pointers
68
int tid = get_thread_num();
69
scalar_t* __restrict__ A = A_tmp + tid * BLOCK_M * K;
70
+ bool use_brgemm = false;
71
72
loop_2d<packed_t>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
73
int64_t n_size = std::min(2 * N - nb * BLOCK_N, BLOCK_N);
@@ -86,6 +84,11 @@ void fused_experts_fp_kernel_impl(
84
bool do_unpack = (mb == mb0) || (expert_id != pre_expert_id);
85
86
int64_t m_size = offsets[mb + 1] - offsets[mb];
87
+ if (m_size == 0) {
88
+ return;
89
+ }
90
+ const bool brg = can_use_brgemm<packed_t>(m_size, n_size);
91
+ use_brgemm |= brg;
92
93
if (nb_offset == 0) {
94
// 1.a load A
@@ -111,7 +114,7 @@ void fused_experts_fp_kernel_impl(
114
/* lda */ K,
115
/* ldb */ n_size,
116
/* ldc */ 2 * N,
114
- /* brg */ use_brgemm,
117
+ /* brg */ brg,
118
/* block_size_K */ block_size_K,
119
/* do_unpack */ do_unpack);
120
});
@@ -152,10 +155,16 @@ void fused_experts_fp_kernel_impl(
155
parallel_2d(MB2, NB2, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
156
int tid = get_thread_num();
157
alignas(64) scalar_t C[BLOCK_M * BLOCK_K];
158
+ bool use_brgemm = false;
159
160
loop_2d<packed_t>(mb0, mb1, nb0, nb1, BLOCK_N * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
161
int64_t m_size = offsets[mb + 1] - offsets[mb];
162
+ if (m_size == 0) {
163
+ return;
164
+ }
165
int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);
166
+ const bool brg = can_use_brgemm<packed_t>(m_size, n_size);
167
+ use_brgemm |= brg;
168
169
// A ptr from ic1 of [M * topk, N] in sorted order
170
// so as to avoid copy A to tmp buffer again
@@ -187,7 +196,7 @@ void fused_experts_fp_kernel_impl(
196
/* lda */ IC,
197
/* ldb */ n_size,
198
/* ldc */ BLOCK_N,
190
- /* brg */ use_brgemm,
199
+ /* brg */ brg,
200
/* block_size_K */ block_size_K,
201
/* do_unpack */ do_unpack);
202
tests/kernels/moe/test_cpu_quant_fused_moe.py
+91
@@ -31,6 +31,21 @@ def _prepack_experts(w: torch.Tensor) -> torch.Tensor:
31
return torch.ops._C.convert_weight_packed(w)
32
33
34
+def _deterministic_expert_routes(
35
+ block_sizes: tuple[int, ...],
36
+) -> tuple[torch.Tensor, torch.Tensor]:
37
+ """Create top-1 routes with exact per-expert block occupancies."""
38
+ topk_ids = torch.cat(
39
+ [
40
+ torch.full((size,), expert, dtype=torch.int32)
41
+ for expert, size in enumerate(block_sizes)
42
+ ]
43
+ )
44
+ topk_ids = topk_ids.view(-1, 1)
45
+ topk_weights = torch.ones(topk_ids.shape, dtype=torch.float32)
46
+ return topk_weights, topk_ids
47
+
48
+
49
# ===========================================================================
50
# FP8 W8A16 MoE
51
# ===========================================================================
@@ -219,6 +234,40 @@ def test_w8a16_block_fp8_cpu_fused_moe(M, N, K, E, topk, seed):
234
torch.testing.assert_close(out_inplace, out, atol=0, rtol=0)
235
236
237
+def test_w8a16_block_fp8_cpu_fused_moe_small_expert_blocks():
238
+ """Test FP8 BRGEMM at exact and multi-block expert boundaries."""
239
+ set_random_seed(0)
240
+ block_sizes = (4, 5, 33)
241
+ N, K, E = 128, 128, len(block_sizes)
242
+ M = sum(block_sizes)
243
+
244
+ a = torch.randn(M, K, dtype=torch.bfloat16) / math.sqrt(K)
245
+ w1, w2, w1_s, w2_s = _make_fp8_moe_weights(E, N, K, BLOCK_SIZE)
246
+ topk_weight, topk_ids = _deterministic_expert_routes(block_sizes)
247
+
248
+ ref_out = ref_w8a16_block_fp8_moe(
249
+ a, w1, w2, w1_s, w2_s, topk_weight, topk_ids, BLOCK_SIZE
250
+ )
251
+ pw1, pw2 = _prepack_experts(w1), _prepack_experts(w2)
252
+ out = ops.fused_experts_cpu(
253
+ a,
254
+ pw1,
255
+ pw2,
256
+ topk_weight,
257
+ topk_ids,
258
+ False,
259
+ ops.CPUQuantMethod.FP8_W8A16,
260
+ w1_s,
261
+ w2_s,
262
+ None,
263
+ None,
264
+ BLOCK_SIZE,
265
+ is_vnni=True,
266
+ )
267
+
268
+ torch.testing.assert_close(ref_out.bfloat16(), out, atol=1e-2, rtol=1e-2)
269
+
270
+
271
# ===========================================================================
272
# MXFP4 W4A16 MoE
273
# ===========================================================================
@@ -438,6 +487,48 @@ def test_mxfp4_cpu_fused_moe(M, N, K, E, topk, seed):
487
torch.testing.assert_close(ref_out.bfloat16(), out, atol=1e-2, rtol=1e-2)
488
489
490
+def test_mxfp4_cpu_fused_moe_small_expert_blocks():
491
+ """Test MXFP4 BRGEMM at exact and multi-block expert boundaries."""
492
+ set_random_seed(0)
493
+ block_sizes = (4, 5, 33)
494
+ N, K, E = 128, 128, len(block_sizes)
495
+ M = sum(block_sizes)
496
+ dtype = torch.bfloat16
497
+
498
+ a = torch.randn(M, K, dtype=dtype) / 10
499
+ w1_bf16 = torch.randn(E, 2 * N, K, dtype=dtype) / 10
500
+ w1q, w1s = MXFP4QuantizeUtil.quantize(w1_bf16)
501
+ w1s = w1s.reshape(E, 2 * N, K // 32)
502
+ w1dq = MXFP4QuantizeUtil.dequantize(w1q, dtype, w1s)
503
+
504
+ w2_bf16 = torch.randn(E, K, N, dtype=dtype) / 10
505
+ w2q, w2s = MXFP4QuantizeUtil.quantize(w2_bf16)
506
+ w2s = w2s.reshape(E, K, N // 32)
507
+ w2dq = MXFP4QuantizeUtil.dequantize(w2q, dtype, w2s)
508
+
509
+ topk_weight, topk_ids = _deterministic_expert_routes(block_sizes)
510
+ ref_out = ref_mxfp4_fused_moe(a, w1dq, w2dq, topk_weight, topk_ids, 1)
511
+
512
+ pw1, pw1s = _prepack_mxfp4_experts(w1q, w1s)
513
+ pw2, pw2s = _prepack_mxfp4_experts(w2q, w2s)
514
+ out = ops.fused_experts_cpu(
515
+ a,
516
+ pw1,
517
+ pw2,
518
+ topk_weight,
519
+ topk_ids,
520
+ False,
521
+ ops.CPUQuantMethod.MXFP4,
522
+ pw1s,
523
+ pw2s,
524
+ None,
525
+ None,
526
+ None,
527
+ )
528
+
529
+ torch.testing.assert_close(ref_out.bfloat16(), out, atol=1e-2, rtol=1e-2)
530
+
531
+
532
@pytest.mark.parametrize("M", [1, 32])
533
@pytest.mark.parametrize("N,K,E,topk", [(128, 128, 4, 2), (64, 64, 4, 2)])
534
@pytest.mark.parametrize("seed", [0])