From bda4c3e8eefe5442897f07f9c0c2d1cae833ca2f Mon Sep 17 00:00:00 2001 From: ccaadaro <90785470+ccaadaro@users.noreply.github.com> Date: Fri, 14 Aug 2026 10:19:21 +0200 Subject: [PATCH] [CPU] Fold the MXFP4 block scale in 2 instructions instead of 4 (#51583) Signed-off-by: ccaadaro Co-authored-by: Li, Jiang --- csrc/cpu/sgl-kernels/vec.h | 27 ++-- tests/kernels/moe/test_cpu_quant_fused_moe.py | 119 ++++++++++++++++++ 2 files changed, 135 insertions(+), 11 deletions(-) diff --git a/csrc/cpu/sgl-kernels/vec.h b/csrc/cpu/sgl-kernels/vec.h index 04b696abf5..c32a01ecd5 100644 --- a/csrc/cpu/sgl-kernels/vec.h +++ b/csrc/cpu/sgl-kernels/vec.h @@ -204,7 +204,6 @@ inline std::tuple<__m512bh, __m512bh> cvt_mxfp4_e2m1_bf16_intrinsic_lut(__m256i const __m512i lut = (__m512i)(_mm512_cvtne2ps_pbh(values, values)); const __m512i abs_mask = _mm512_set1_epi16(0x7FFF); - const __m512i zero = _mm512_setzero_si512(); // expand values to 16-bit integers __m512i x0 = _mm512_cvtepu8_epi16(a); @@ -214,17 +213,23 @@ inline std::tuple<__m512bh, __m512bh> cvt_mxfp4_e2m1_bf16_intrinsic_lut(__m256i x0 = _mm512_permutexvar_epi16(x0, lut); x1 = _mm512_permutexvar_epi16(x1, lut); - // check for zeros - __mmask32 mask0 = _mm512_cmp_epi16_mask(_mm512_and_si512(x0, abs_mask), zero, _MM_CMPINT_EQ); - __mmask32 mask1 = _mm512_cmp_epi16_mask(_mm512_and_si512(x1, abs_mask), zero, _MM_CMPINT_EQ); + // Emulate the bf16 multiply by the E8M0 scale as an integer add on the + // exponent field, forcing zeros to stay zero. + // + // vptestmw sets a lane's mask bit when (x & 0x7FFF) != 0, i.e. for the lanes + // that are not +0.0 or -0.0. That is the *complement* of the mask the old + // and/cmp/add/blend sequence built, which selected the zero lanes; the + // zero-masking form of vpaddw then writes 0 where the mask is clear, so the + // two agree lane by lane. That includes the LUT's -0.0 entry, which both + // versions turn into +0.0. + // + // 2 instructions per vector instead of 4, no extra ISA requirement: + // vptestmw is AVX512BW, which vpermw above already needs. + __mmask32 mask0 = _mm512_test_epi16_mask(x0, abs_mask); + __mmask32 mask1 = _mm512_test_epi16_mask(x1, abs_mask); - // emulate bf16 mul with scale factor - x0 = _mm512_add_epi16(x0, s0); - x1 = _mm512_add_epi16(x1, s1); - - // blend with zero - x0 = _mm512_mask_blend_epi16(mask0, x0, zero); - x1 = _mm512_mask_blend_epi16(mask1, x1, zero); + x0 = _mm512_maskz_add_epi16(mask0, x0, s0); + x1 = _mm512_maskz_add_epi16(mask1, x1, s1); return std::make_tuple(__m512bh(x0), __m512bh(x1)); } diff --git a/tests/kernels/moe/test_cpu_quant_fused_moe.py b/tests/kernels/moe/test_cpu_quant_fused_moe.py index 61e8aeb7a9..e039310b93 100644 --- a/tests/kernels/moe/test_cpu_quant_fused_moe.py +++ b/tests/kernels/moe/test_cpu_quant_fused_moe.py @@ -529,6 +529,125 @@ def test_mxfp4_cpu_fused_moe_small_expert_blocks(): torch.testing.assert_close(ref_out.bfloat16(), out, atol=1e-2, rtol=1e-2) +# Both E2M1 zero codes: 0b0000 (+0.0) and 0b1000 (-0.0), in both nibbles. +MXFP4_ZERO_BYTES = [0x00, 0x88, 0x08, 0x80] +# 0 and 255 are the ends of the E8M0 range; 127 is the identity scale. +MXFP4_E8M0_VALUES = [0, 1, 127, 200, 254, 255] + + +@pytest.mark.parametrize("zero_byte", MXFP4_ZERO_BYTES) +@pytest.mark.parametrize("e8m0", MXFP4_E8M0_VALUES) +def test_mxfp4_cpu_zero_codes_stay_zero(zero_byte, e8m0): + """A zero E2M1 code stays zero for every E8M0 exponent. + + The unpack applies the block scale as an integer add on the bf16 exponent + field, which is exact for every value in the E2M1 codebook except the two + zeros: 0x0000 and 0x8000 have no exponent to shift, so adding to them + produces a small finite number instead of zero. Both are special-cased, and + this is the invariant that special case exists for. + + Worth pinning separately from ``test_mxfp4_cpu_fused_moe``: there the zero + codes are a small fraction of random weights and a broken special case + would stay inside the 1e-2 tolerance for the low exponents. Here every + weight is a zero, so the output is exactly zero or it is not. + """ + N, K, E, M = 64, 64, 2, 4 + dtype = torch.bfloat16 + set_random_seed(0) + + a = torch.randn(M, K, dtype=dtype) + w1q = torch.full((E, 2 * N, K // 2), zero_byte, dtype=torch.uint8) + w1s = torch.full((E, 2 * N, K // 32), e8m0, dtype=torch.uint8) + # w2 is ordinary: the zeros have to survive the first GEMM and the + # activation, and a nonzero w2 is what would expose it if they did not. + w2_bf16 = torch.randn(E, K, N, dtype=dtype) / 10 + w2q, w2s = MXFP4QuantizeUtil.quantize(w2_bf16) + w2s = w2s.reshape(E, K, N // 32) + + topk_weight = torch.ones((M, 1), dtype=torch.float32) + topk_ids = torch.zeros((M, 1), dtype=torch.int32) + + pw1, pw1s = _prepack_mxfp4_experts(w1q, w1s) + pw2, pw2s = _prepack_mxfp4_experts(w2q, w2s) + out = ops.fused_experts_cpu( + a, + pw1, + pw2, + topk_weight, + topk_ids, + False, + ops.CPUQuantMethod.MXFP4, + pw1s, + pw2s, + None, + None, + None, + ) + + # silu(0) * 0 = 0, so the whole layer collapses to exactly zero. Not + # assert_close: any nonzero output here is a wrong unpack, not rounding. + assert torch.equal(out, torch.zeros_like(out)), ( + f"zero code 0x{zero_byte:02x} with e8m0={e8m0} produced " + f"max |out| = {out.abs().max().item()}" + ) + + +# Narrower than MXFP4_E8M0_VALUES on purpose: this test keeps nonzero weights, +# and 6.0 * 2**(255-127) is not representable in bf16, so the ends of the E8M0 +# range would compare inf against inf and prove nothing. The all-zero test above +# is the one that can reach them. +MXFP4_E8M0_FINITE = [107, 127, 137] + + +@pytest.mark.parametrize("e8m0", MXFP4_E8M0_FINITE) +def test_mxfp4_cpu_zero_codes_mixed_with_nonzero(e8m0): + """Zeros and nonzeros in the same 32-element scale block. + + The zero check is per lane, not per block: this fails if it is ever + rewritten as a whole-block branch. Uses a single shared exponent so the + reference is a plain power of two. + """ + N, K, E, M = 64, 64, 2, 4 + dtype = torch.bfloat16 + set_random_seed(0) + + a = torch.randn(M, K, dtype=dtype) + # Alternate a zero byte and a nonzero one along K, so every scale block + # holds both kinds. + pattern = torch.tensor([0x88, 0x21], dtype=torch.uint8).repeat(K // 4) + w1q = pattern.view(1, 1, -1).expand(E, 2 * N, K // 2).contiguous() + w1s = torch.full((E, 2 * N, K // 32), e8m0, dtype=torch.uint8) + w1dq = MXFP4QuantizeUtil.dequantize(w1q, dtype, w1s) + + w2_bf16 = torch.randn(E, K, N, dtype=dtype) / 10 + w2q, w2s = MXFP4QuantizeUtil.quantize(w2_bf16) + w2s = w2s.reshape(E, K, N // 32) + w2dq = MXFP4QuantizeUtil.dequantize(w2q, dtype, w2s) + + topk_weight = torch.ones((M, 1), dtype=torch.float32) + topk_ids = torch.zeros((M, 1), dtype=torch.int32) + ref_out = ref_mxfp4_fused_moe(a, w1dq, w2dq, topk_weight, topk_ids, 1) + + pw1, pw1s = _prepack_mxfp4_experts(w1q, w1s) + pw2, pw2s = _prepack_mxfp4_experts(w2q, w2s) + out = ops.fused_experts_cpu( + a, + pw1, + pw2, + topk_weight, + topk_ids, + False, + ops.CPUQuantMethod.MXFP4, + pw1s, + pw2s, + None, + None, + None, + ) + + torch.testing.assert_close(ref_out.bfloat16(), out, atol=1e-2, rtol=1e-2) + + @pytest.mark.parametrize("M", [1, 32]) @pytest.mark.parametrize("N,K,E,topk", [(128, 128, 4, 2), (64, 64, 4, 2)]) @pytest.mark.parametrize("seed", [0])