[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>
This commit is contained in:
ccaadaro
co-authored by
Li, Jiang <jiang1.li@intel.com>
Li, Jiang <jiang1.li@intel.com>
parent
8e6d8e4f6a
commit
bda4c3e8ee
+16
-11
@@ -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));
|
||||
}
|
||||
|
||||
@@ -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])
|
||||
|
||||
Reference in New Issue
Block a user