[CPU] Fold the MXFP4 block scale in 2 instructions instead of 4 (#51583)
Signed-off-by: ccaadaro <ccanadar@unex.es> Co-authored-by: Li, Jiang <jiang1.li@intel.com>
ccaadaro committed
Aug 14, 2026 at 10:19 UTC
bda4c3e8eefe5442897f07f9c0c2d1cae833ca2f
2 files changed
+135
-11
csrc/cpu/sgl-kernels/vec.h
+16
-11
@@ -204,7 +204,6 @@ inline std::tuple<__m512bh, __m512bh> cvt_mxfp4_e2m1_bf16_intrinsic_lut(__m256i
204
const __m512i lut = (__m512i)(_mm512_cvtne2ps_pbh(values, values));
205
206
const __m512i abs_mask = _mm512_set1_epi16(0x7FFF);
207
- const __m512i zero = _mm512_setzero_si512();
207
208
// expand values to 16-bit integers
209
__m512i x0 = _mm512_cvtepu8_epi16(a);
@@ -214,17 +213,23 @@ inline std::tuple<__m512bh, __m512bh> cvt_mxfp4_e2m1_bf16_intrinsic_lut(__m256i
213
x0 = _mm512_permutexvar_epi16(x0, lut);
214
x1 = _mm512_permutexvar_epi16(x1, lut);
215
217
- // check for zeros
218
- __mmask32 mask0 = _mm512_cmp_epi16_mask(_mm512_and_si512(x0, abs_mask), zero, _MM_CMPINT_EQ);
219
- __mmask32 mask1 = _mm512_cmp_epi16_mask(_mm512_and_si512(x1, abs_mask), zero, _MM_CMPINT_EQ);
220
-
221
- // emulate bf16 mul with scale factor
222
- x0 = _mm512_add_epi16(x0, s0);
223
- x1 = _mm512_add_epi16(x1, s1);
216
+ // Emulate the bf16 multiply by the E8M0 scale as an integer add on the
217
+ // exponent field, forcing zeros to stay zero.
218
+ //
219
+ // vptestmw sets a lane's mask bit when (x & 0x7FFF) != 0, i.e. for the lanes
220
+ // that are not +0.0 or -0.0. That is the *complement* of the mask the old
221
+ // and/cmp/add/blend sequence built, which selected the zero lanes; the
222
+ // zero-masking form of vpaddw then writes 0 where the mask is clear, so the
223
+ // two agree lane by lane. That includes the LUT's -0.0 entry, which both
224
+ // versions turn into +0.0.
225
+ //
226
+ // 2 instructions per vector instead of 4, no extra ISA requirement:
227
+ // vptestmw is AVX512BW, which vpermw above already needs.
228
+ __mmask32 mask0 = _mm512_test_epi16_mask(x0, abs_mask);
229
+ __mmask32 mask1 = _mm512_test_epi16_mask(x1, abs_mask);
230
225
- // blend with zero
226
- x0 = _mm512_mask_blend_epi16(mask0, x0, zero);
227
- x1 = _mm512_mask_blend_epi16(mask1, x1, zero);
231
+ x0 = _mm512_maskz_add_epi16(mask0, x0, s0);
232
+ x1 = _mm512_maskz_add_epi16(mask1, x1, s1);
233
234
return std::make_tuple(__m512bh(x0), __m512bh(x1));
235
}
tests/kernels/moe/test_cpu_quant_fused_moe.py
+119
@@ -529,6 +529,125 @@ def test_mxfp4_cpu_fused_moe_small_expert_blocks():
529
torch.testing.assert_close(ref_out.bfloat16(), out, atol=1e-2, rtol=1e-2)
530
531
532
+# Both E2M1 zero codes: 0b0000 (+0.0) and 0b1000 (-0.0), in both nibbles.
533
+MXFP4_ZERO_BYTES = [0x00, 0x88, 0x08, 0x80]
534
+# 0 and 255 are the ends of the E8M0 range; 127 is the identity scale.
535
+MXFP4_E8M0_VALUES = [0, 1, 127, 200, 254, 255]
536
+
537
+
538
+@pytest.mark.parametrize("zero_byte", MXFP4_ZERO_BYTES)
539
+@pytest.mark.parametrize("e8m0", MXFP4_E8M0_VALUES)
540
+def test_mxfp4_cpu_zero_codes_stay_zero(zero_byte, e8m0):
541
+ """A zero E2M1 code stays zero for every E8M0 exponent.
542
+
543
+ The unpack applies the block scale as an integer add on the bf16 exponent
544
+ field, which is exact for every value in the E2M1 codebook except the two
545
+ zeros: 0x0000 and 0x8000 have no exponent to shift, so adding to them
546
+ produces a small finite number instead of zero. Both are special-cased, and
547
+ this is the invariant that special case exists for.
548
+
549
+ Worth pinning separately from ``test_mxfp4_cpu_fused_moe``: there the zero
550
+ codes are a small fraction of random weights and a broken special case
551
+ would stay inside the 1e-2 tolerance for the low exponents. Here every
552
+ weight is a zero, so the output is exactly zero or it is not.
553
+ """
554
+ N, K, E, M = 64, 64, 2, 4
555
+ dtype = torch.bfloat16
556
+ set_random_seed(0)
557
+
558
+ a = torch.randn(M, K, dtype=dtype)
559
+ w1q = torch.full((E, 2 * N, K // 2), zero_byte, dtype=torch.uint8)
560
+ w1s = torch.full((E, 2 * N, K // 32), e8m0, dtype=torch.uint8)
561
+ # w2 is ordinary: the zeros have to survive the first GEMM and the
562
+ # activation, and a nonzero w2 is what would expose it if they did not.
563
+ w2_bf16 = torch.randn(E, K, N, dtype=dtype) / 10
564
+ w2q, w2s = MXFP4QuantizeUtil.quantize(w2_bf16)
565
+ w2s = w2s.reshape(E, K, N // 32)
566
+
567
+ topk_weight = torch.ones((M, 1), dtype=torch.float32)
568
+ topk_ids = torch.zeros((M, 1), dtype=torch.int32)
569
+
570
+ pw1, pw1s = _prepack_mxfp4_experts(w1q, w1s)
571
+ pw2, pw2s = _prepack_mxfp4_experts(w2q, w2s)
572
+ out = ops.fused_experts_cpu(
573
+ a,
574
+ pw1,
575
+ pw2,
576
+ topk_weight,
577
+ topk_ids,
578
+ False,
579
+ ops.CPUQuantMethod.MXFP4,
580
+ pw1s,
581
+ pw2s,
582
+ None,
583
+ None,
584
+ None,
585
+ )
586
+
587
+ # silu(0) * 0 = 0, so the whole layer collapses to exactly zero. Not
588
+ # assert_close: any nonzero output here is a wrong unpack, not rounding.
589
+ assert torch.equal(out, torch.zeros_like(out)), (
590
+ f"zero code 0x{zero_byte:02x} with e8m0={e8m0} produced "
591
+ f"max |out| = {out.abs().max().item()}"
592
+ )
593
+
594
+
595
+# Narrower than MXFP4_E8M0_VALUES on purpose: this test keeps nonzero weights,
596
+# and 6.0 * 2**(255-127) is not representable in bf16, so the ends of the E8M0
597
+# range would compare inf against inf and prove nothing. The all-zero test above
598
+# is the one that can reach them.
599
+MXFP4_E8M0_FINITE = [107, 127, 137]
600
+
601
+
602
+@pytest.mark.parametrize("e8m0", MXFP4_E8M0_FINITE)
603
+def test_mxfp4_cpu_zero_codes_mixed_with_nonzero(e8m0):
604
+ """Zeros and nonzeros in the same 32-element scale block.
605
+
606
+ The zero check is per lane, not per block: this fails if it is ever
607
+ rewritten as a whole-block branch. Uses a single shared exponent so the
608
+ reference is a plain power of two.
609
+ """
610
+ N, K, E, M = 64, 64, 2, 4
611
+ dtype = torch.bfloat16
612
+ set_random_seed(0)
613
+
614
+ a = torch.randn(M, K, dtype=dtype)
615
+ # Alternate a zero byte and a nonzero one along K, so every scale block
616
+ # holds both kinds.
617
+ pattern = torch.tensor([0x88, 0x21], dtype=torch.uint8).repeat(K // 4)
618
+ w1q = pattern.view(1, 1, -1).expand(E, 2 * N, K // 2).contiguous()
619
+ w1s = torch.full((E, 2 * N, K // 32), e8m0, dtype=torch.uint8)
620
+ w1dq = MXFP4QuantizeUtil.dequantize(w1q, dtype, w1s)
621
+
622
+ w2_bf16 = torch.randn(E, K, N, dtype=dtype) / 10
623
+ w2q, w2s = MXFP4QuantizeUtil.quantize(w2_bf16)
624
+ w2s = w2s.reshape(E, K, N // 32)
625
+ w2dq = MXFP4QuantizeUtil.dequantize(w2q, dtype, w2s)
626
+
627
+ topk_weight = torch.ones((M, 1), dtype=torch.float32)
628
+ topk_ids = torch.zeros((M, 1), dtype=torch.int32)
629
+ ref_out = ref_mxfp4_fused_moe(a, w1dq, w2dq, topk_weight, topk_ids, 1)
630
+
631
+ pw1, pw1s = _prepack_mxfp4_experts(w1q, w1s)
632
+ pw2, pw2s = _prepack_mxfp4_experts(w2q, w2s)
633
+ out = ops.fused_experts_cpu(
634
+ a,
635
+ pw1,
636
+ pw2,
637
+ topk_weight,
638
+ topk_ids,
639
+ False,
640
+ ops.CPUQuantMethod.MXFP4,
641
+ pw1s,
642
+ pw2s,
643
+ None,
644
+ None,
645
+ None,
646
+ )
647
+
648
+ torch.testing.assert_close(ref_out.bfloat16(), out, atol=1e-2, rtol=1e-2)
649
+
650
+
651
@pytest.mark.parametrize("M", [1, 32])
652
@pytest.mark.parametrize("N,K,E,topk", [(128, 128, 4, 2), (64, 64, 4, 2)])
653
@pytest.mark.parametrize("seed", [0])