diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/nvfp4.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/nvfp4.cpp index 3c5a56d51f..ffd9d788e9 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/kernels/nvfp4.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/nvfp4.cpp @@ -12,10 +12,11 @@ namespace { // Activations stay bf16 (W4A16); the GEMV dequantizes weights inside the K-loop. -// e2m1 * 2 as exact int8 (all e2m1 values are multiples of 0.5). Used by the AVX-VNNI -// W4A8 path: nibble -> int8 weight via PSHUFB LUT, then VPDPBUSD against int8 activations; -// the *2 is undone by a 0.5 folded into the final scale. Mirrors ggml's kvalues_mxfp4. +// e2m1 * 2 as exact int8 (all e2m1 values are multiples of 0.5); the W4A8 dots undo the *2 with a 0.5 in the final scale. Mirrors ggml's kvalues_mxfp4. alignas(16) const int8_t kE2M1x2[16] = {0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12}; +// kE2M1x2 + kE2M1Bias, the PSHUFB LUT of the VNNI dots: VPDPBUSD takes the weight as the unsigned operand, and W4a8Act::corr carries the -kE2M1Bias * sum(a) correction. +alignas(16) const uint8_t kE2M1x2u[16] = {12, 13, 14, 15, 16, 18, 20, 24, 12, 11, 10, 9, 8, 6, 4, 0}; +constexpr int kE2M1Bias = 12; // Activations pre-deinterleaved to fp32 (xe[m]=x[2m], xo[m]=x[2m+1]); see the // ds_fp4 dot in ds_fp4.cpp for why (drops the hot loop to ~1.5 shuffle ops / 16 weights). @@ -44,29 +45,16 @@ float dot_nvfp4_scalar(const uint8_t* packed, const uint8_t* scale, float global } // ---- NVFP4 W4A8 (int8 activations) dot: nibble->int8 via LUT, per-16 act scale ---- -// asi8: int8 activations laid out per-16 block as [even(8), odd(8)]; asb[b] = per-block -// activation scale (absmax/127). Result folds the e2m1*2 -> *0.5 into the scale. -using nvi8dot_fn = float (*)(const uint8_t*, const uint8_t*, float, const int8_t*, int, - const float*, const float*); - -[[maybe_unused]] float dot_nvfp4_i8_scalar(const uint8_t* packed, const uint8_t* scale, - float global, const int8_t* asi8, int K, const float* e4m3, - const float* asb) { - float acc = 0.0f; - const int nb = K / 16; - for (int b = 0; b < nb; ++b) { - const uint8_t* pk = packed + (size_t)b * 8; - const int8_t* ae = asi8 + (size_t)b * 16; // even(8) - const int8_t* ao = ae + 8; // odd(8) - int isum = 0; - for (int j = 0; j < 8; ++j) { - isum += (int)kE2M1x2[pk[j] & 0xF] * (int)ae[j]; - isum += (int)kE2M1x2[pk[j] >> 4] * (int)ao[j]; - } - acc += (e4m3[scale[b]] * asb[b]) * (float)isum; - } - return acc * (0.5f * global); -} +// Activations are grouped in chunks of 8 16-K blocks: the even-K int8 of the chunk's blocks, then its odd-K int8. +// A 64-byte load of packed weights then splits into low (even-K) and high (odd-K) nibbles that line up with the activations, and int32 lane l of a dot holds block l / 2. +struct W4a8Act { + const int8_t* q; // [K] + const float* scale; // [K/16], absmax / 127 of each block + const int32_t* corr; // [K/8], one per int32 lane: -kE2M1Bias * the sum of the 8 activations it multiplies +}; + +using nvi8dot_fn = float (*)(const uint8_t*, const uint8_t*, float, const W4a8Act&, int); +using nvi8quant_fn = void (*)(const bf16_t*, int, int8_t*, float*, int32_t*); #if CPU_MOE_X86 // Two 16-K blocks (16 packed bytes) per iter: lo nibbles -> even-K, hi -> odd-K, @@ -147,54 +135,6 @@ float dot_nvfp4_avx2(const uint8_t* packed, const uint8_t* scale, float global, return hsum256(_mm256_add_ps(acc0, acc1)) * global; } -// AVX-VNNI W4A8: decode 8 packed bytes (16 nibbles) of one 16-block to int8 [lo(8),hi(8)] -// via PSHUFB against the e2m1*2 LUT (replaces the 2 vpermps fp32 expands -- ~4x less -// port-5 traffic). lo=even-K weights, hi=odd-K, matching the [even(8),odd(8)] act layout. -CPU_MOE_TARGET("avx2,avxvnni,fma") -inline __m128i nvfp4_decode_block_i8(const uint8_t* pk, __m128i lut) { - __m128i b = _mm_loadl_epi64(reinterpret_cast(pk)); // 8 bytes - __m128i lo = _mm_and_si128(b, _mm_set1_epi8(0x0F)); - __m128i hi = _mm_and_si128(_mm_srli_epi16(b, 4), _mm_set1_epi8(0x0F)); - return _mm_shuffle_epi8(lut, _mm_unpacklo_epi64(lo, hi)); // [lo(8),hi(8)] -> int8 -} - -// Two 16-blocks per VPDPBUSD (32 int8). Sign trick (ggml mul_add_epi8): |w|*(sign(w)*a)=w*a, -// so u8*s8 needs no offset/correction term. Per-block scale (e4m3 * act-scale) folded in fp32 -// (lanes 0-3 -> block b, 4-7 -> block b+1). Bit-faithful weight; only the int8 activation quant -// (W4A8) differs from the bf16 reference. -CPU_MOE_TARGET("avx2,avxvnni,fma") -float dot_nvfp4_i8_vnni(const uint8_t* packed, const uint8_t* scale, float global, - const int8_t* asi8, int K, const float* e4m3, const float* asb) { - const __m128i lut = _mm_loadu_si128(reinterpret_cast(kE2M1x2)); - __m256 accF = _mm256_setzero_ps(); - const int nb = K / 16; - int b = 0; - for (; b + 2 <= nb; b += 2) { - __m128i wb = nvfp4_decode_block_i8(packed + (size_t)b * 8, lut); - __m128i wb1 = nvfp4_decode_block_i8(packed + (size_t)(b + 1) * 8, lut); - __m256i w = _mm256_set_m128i(wb1, wb); // [blk b | blk b+1] - __m256i a = _mm256_loadu_si256(reinterpret_cast(asi8 + (size_t)b * 16)); - __m256i aw = _mm256_sign_epi8(w, w); // |w| (u8 operand) - __m256i sa = _mm256_sign_epi8(a, w); // sign(w)*a (s8 operand) - __m256i di = _mm256_dpbusd_avx_epi32(_mm256_setzero_si256(), aw, sa); - __m256 scv = _mm256_blend_ps(_mm256_set1_ps(e4m3[scale[b]] * asb[b]), - _mm256_set1_ps(e4m3[scale[b + 1]] * asb[b + 1]), 0xF0); - accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(di), scv, accF); - } - float s = hsum256(accF); - for (; b < nb; ++b) { // tail (odd block count) - const uint8_t* pk = packed + (size_t)b * 8; - const int8_t* ae = asi8 + (size_t)b * 16; const int8_t* ao = ae + 8; - int isum = 0; - for (int j = 0; j < 8; ++j) - isum += (int)kE2M1x2[pk[j] & 0xF] * (int)ae[j] + (int)kE2M1x2[pk[j] >> 4] * (int)ao[j]; - s += (e4m3[scale[b]] * asb[b]) * (float)isum; - } - return s * (0.5f * global); -} - -#ifdef CPU_MOE_HAS_AVX512VNNI - // Software-prefetch distance for the W4A8 weight stream, in 16-K blocks (8 packed // bytes each). Returns -1 when FREETOKEN_CPU_MOE_PF_BLOCKS is unset: the kernel then // uses the built-in default min(512 blocks = 4 KB, 2 rows) -- 4 KB is the empirical @@ -213,96 +153,239 @@ static int nvfp4_pf_blocks() { return v; } -// AVX-512 VNNI W4A8: FOUR 16-K blocks per VPDPBUSD (64 int8) -- 2x the AVX-VNNI -// (256-bit) path. Decode 32 packed bytes -> 64 int8 with a single _mm512_shuffle_epi8 -// (e2m1*2 LUT replicated to all 4 128-bit lanes). AVX-512 has no _mm512_sign_epi8, so -// the u8*s8 sign trick (|w| as u8, sign(w)*a as s8) uses abs_epi8 + a masked negate. -// Bit-faithful weight; only the int8 activation quant (W4A8) differs from bf16. -// One 4-block group (64 int8) -> scaled fp32 partial (16 lanes). Isolated as a helper so -// the caller can run several independent chains into separate accumulators (the decode -> -// dpbusd -> scale chain is long, so a single accumulator leaves the core latency-bound). -CPU_MOE_TARGET("avx512f,avx512bw,avx512vnni,avx2") -static inline __m512 nvfp4_i8_grp4(const uint8_t* packed, const uint8_t* scale, - const int8_t* asi8, const float* e4m3, const float* asb, - int b, __m512i lut, __m512i idx, __m512i mask0F, - __m512i idxsc) { - const __mmask64 hi_half = 0xFF00FF00FF00FF00ULL; // bytes 8-15 of each 128b lane -> hi nibbles - __m256i raw = _mm256_loadu_si256(reinterpret_cast(packed + (size_t)b * 8)); - __m512i src = _mm512_permutexvar_epi64(idx, _mm512_castsi256_si512(raw)); - __m512i lo = _mm512_and_si512(src, mask0F); - __m512i hi = _mm512_and_si512(_mm512_srli_epi16(src, 4), mask0F); - __m512i comb = _mm512_mask_blend_epi8(hi_half, lo, hi); - __m512i w = _mm512_shuffle_epi8(lut, comb); // 64 int8 weights (e2m1*2) - // u8*s8 sign trick without _mm512_sign_epi8: aw=|w|, sa = (w<0 ? -a : a) (a moot at w==0) - __m512i a = _mm512_loadu_si512(reinterpret_cast(asi8 + (size_t)b * 16)); - __m512i aw = _mm512_abs_epi8(w); - __mmask64 neg = _mm512_movepi8_mask(w); - __m512i sa = _mm512_mask_sub_epi8(a, neg, _mm512_setzero_si512(), a); - __m512i di = _mm512_dpbusd_epi32(_mm512_setzero_si512(), aw, sa); // 16 int32 (groups of 4) - // int32[0..3]->blk b, [4..7]->b+1, [8..11]->b+2, [12..15]->b+3. The 4 block scales - // (e4m3 LUT x per-block act scale) are computed vectorized: a 4-byte load + epu8->epi32 - // widen + one 4-lane LUT gather + one mul replaces 8 scalar loads + 4 scalar muls + - // a set_ps assembly, which otherwise dominates the per-group op count. - int sc_raw; - memcpy(&sc_raw, scale + b, 4); - __m128i sc4 = _mm_cvtepu8_epi32(_mm_cvtsi32_si128(sc_raw)); - __m128 s4 = _mm_mul_ps(_mm_i32gather_ps(e4m3, sc4, 4), _mm_loadu_ps(asb + b)); - __m512 scv = _mm512_permutexvar_ps(idxsc, _mm512_castps128_ps512(s4)); - return _mm512_mul_ps(_mm512_cvtepi32_ps(di), scv); +// e4m3 -> fp16 is a shift: exponent and mantissa move up 7 bits, and the bias gap (15 - 7) leaves a 2^-8 that the kernels undo in the final scale. +// The byte is sign-extended first, so its sign also lands in bit 15 and & 0xBF80 keeps it. +constexpr float kE4m3ViaF16 = 256.0f; +constexpr int16_t kE4m3F16Mask = (int16_t)0xBF80; + +int w4a8_pf(int nb) { + const int pfb = nvfp4_pf_blocks(); + return (pfb < 0) ? std::min(512, 2 * nb) : pfb; } -CPU_MOE_TARGET("avx512f,avx512bw,avx512vnni,avx2") -float dot_nvfp4_i8_avx512vnni(const uint8_t* packed, const uint8_t* scale, float global, - const int8_t* asi8, int K, const float* e4m3, const float* asb) { - const __m512i lut = _mm512_broadcast_i32x4( - _mm_loadu_si128(reinterpret_cast(kE2M1x2))); - const __m512i idx = _mm512_set_epi64(3, 3, 2, 2, 1, 1, 0, 0); // block i -> 128b lane i - const __m512i mask0F = _mm512_set1_epi8(0x0F); - const __m512i idxsc = _mm512_set_epi32(3, 3, 3, 3, 2, 2, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0); - __m512 acc0 = _mm512_setzero_ps(), acc1 = _mm512_setzero_ps(); - __m512 acc2 = _mm512_setzero_ps(), acc3 = _mm512_setzero_ps(); +// Weight scale * act scale of 8 blocks, in 2^-8 units. +CPU_MOE_TARGET("avx2,fma,f16c") +inline __m256 w4a8_scales256(__m128i s8, const float* asc) { + __m128i h = _mm_cvtepi8_epi16(s8); + h = _mm_and_si128(_mm_slli_epi16(h, 7), _mm_set1_epi16(kE4m3F16Mask)); + return _mm256_mul_ps(_mm256_cvtph_ps(h), _mm256_loadu_ps(asc)); +} + +CPU_MOE_TARGET("avx2,fma,f16c") +inline void w4a8_split256(__m256i v, __m256i lut, __m256i& lo, __m256i& hi) { + const __m256i m4 = _mm256_set1_epi8(0x0F); + lo = _mm256_shuffle_epi8(lut, _mm256_and_si256(v, m4)); + hi = _mm256_shuffle_epi8(lut, _mm256_and_si256(_mm256_srli_epi16(v, 4), m4)); +} + +CPU_MOE_TARGET("avx2,fma,f16c,avxvnni") +inline __m256i w4a8_dot256_vnni(__m256i corr, __m256i lo, __m256i qe, __m256i hi, __m256i qo) { + return _mm256_dpbusd_avx_epi32(_mm256_dpbusd_avx_epi32(corr, lo, qe), hi, qo); +} + +// sum of e4m3(scale) * act scale * int dot over blocks [b0, nb) in 2^-8 units, before the final 0.5 * global +CPU_MOE_TARGET("avx2,fma,f16c") +float w4a8_blocks_scalar(const uint8_t* packed, const uint8_t* scale, const W4a8Act& a, int b0, int nb) { + float acc = 0.0f; + for (int b = b0; b < nb; ++b) { + const int c0 = b & ~7, nbc = std::min(8, nb - c0); + const int8_t* qe = a.q + (size_t)c0 * 16 + (size_t)(b - c0) * 8; + const int8_t* qo = qe + nbc * 8; + const uint8_t* pk = packed + (size_t)b * 8; + int isum = 0; + for (int j = 0; j < 8; ++j) isum += kE2M1x2[pk[j] & 0xF] * qe[j] + kE2M1x2[pk[j] >> 4] * qo[j]; + acc += (_cvtsh_ss((uint16_t)(((int16_t)(int8_t)scale[b] << 7) & kE4m3F16Mask)) * a.scale[b]) * (float)isum; + } + return acc; +} + +CPU_MOE_TARGET("avx2,fma,f16c,avxvnni") +float dot_nvfp4_i8_vnni(const uint8_t* packed, const uint8_t* scale, float global, const W4a8Act& a, int K) { + const __m256i lut = _mm256_broadcastsi128_si256(_mm_load_si128(reinterpret_cast(kE2M1x2u))); + const __m256i lanes_lo = _mm256_set_epi32(3, 3, 2, 2, 1, 1, 0, 0); + const __m256i lanes_hi = _mm256_set_epi32(7, 7, 6, 6, 5, 5, 4, 4); + const int nb = K / 16, pf = w4a8_pf(nb); + __m256 acc0 = _mm256_setzero_ps(), acc1 = _mm256_setzero_ps(); + __m256i lo, hi; + int b = 0; + for (; b + 8 <= nb; b += 8) { + const uint8_t* pk = packed + (size_t)b * 8; + if (pf > 0) _mm_prefetch(reinterpret_cast(pk + (size_t)pf * 8), _MM_HINT_T0); + const int8_t* q = a.q + (size_t)b * 16; + const int32_t* c = a.corr + (size_t)b * 2; + const __m256 sc = w4a8_scales256(_mm_loadl_epi64(reinterpret_cast(scale + b)), a.scale + b); + w4a8_split256(_mm256_loadu_si256(reinterpret_cast(pk)), lut, lo, hi); + __m256i d = w4a8_dot256_vnni(_mm256_loadu_si256(reinterpret_cast(c)), lo, + _mm256_loadu_si256(reinterpret_cast(q)), hi, + _mm256_loadu_si256(reinterpret_cast(q + 64))); + acc0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d), _mm256_permutevar8x32_ps(sc, lanes_lo), acc0); + w4a8_split256(_mm256_loadu_si256(reinterpret_cast(pk + 32)), lut, lo, hi); + d = w4a8_dot256_vnni(_mm256_loadu_si256(reinterpret_cast(c + 8)), lo, + _mm256_loadu_si256(reinterpret_cast(q + 32)), hi, + _mm256_loadu_si256(reinterpret_cast(q + 96))); + acc1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d), _mm256_permutevar8x32_ps(sc, lanes_hi), acc1); + } + // A last chunk of 4..7 blocks runs its first half here (its odd-K bytes start 8 * n after the even-K ones); fewer than 4 left go scalar. + if (b + 4 <= nb) { + const int n = nb - b; + const int8_t* q = a.q + (size_t)b * 16; + int s4; + std::memcpy(&s4, scale + b, 4); + float as8[8] = {a.scale[b], a.scale[b + 1], a.scale[b + 2], a.scale[b + 3]}; + const __m256 sc = w4a8_scales256(_mm_cvtsi32_si128(s4), as8); + w4a8_split256(_mm256_loadu_si256(reinterpret_cast(packed + (size_t)b * 8)), lut, lo, hi); + const __m256i d = w4a8_dot256_vnni(_mm256_loadu_si256(reinterpret_cast(a.corr + (size_t)b * 2)), lo, + _mm256_loadu_si256(reinterpret_cast(q)), hi, + _mm256_loadu_si256(reinterpret_cast(q + 8 * n))); + acc0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d), _mm256_permutevar8x32_ps(sc, lanes_lo), acc0); + b += 4; + } + const float tail = w4a8_blocks_scalar(packed, scale, a, b, nb); + return (hsum256(_mm256_add_ps(acc0, acc1)) + tail) * (0.5f * kE4m3ViaF16 * global); +} + +CPU_MOE_TARGET("avx2") +inline int hsum_epi32_256(__m256i v) { + __m128i s = _mm_add_epi32(_mm256_castsi256_si128(v), _mm256_extracti128_si256(v, 1)); + s = _mm_add_epi32(s, _mm_shuffle_epi32(s, 0x4E)); + s = _mm_add_epi32(s, _mm_shuffle_epi32(s, 0xB1)); + return _mm_cvtsi128_si32(s); +} + +CPU_MOE_TARGET("avx2,fma") +inline __m256i w4a8_round256(__m256 y) { + const __m256 absm = _mm256_castsi256_ps(_mm256_set1_epi32(0x7FFFFFFF)); + const __m256 t = _mm256_round_ps(y, _MM_FROUND_TO_ZERO | _MM_FROUND_NO_EXC); + const __m256 away = _mm256_cmp_ps(_mm256_and_ps(_mm256_sub_ps(y, t), absm), _mm256_set1_ps(0.5f), _CMP_GE_OQ); + const __m256 step = _mm256_or_ps(_mm256_and_ps(y, _mm256_set1_ps(-0.0f)), _mm256_set1_ps(1.0f)); + const __m256 r = _mm256_add_ps(t, _mm256_and_ps(away, step)); + return _mm256_cvtps_epi32(_mm256_min_ps(_mm256_max_ps(r, _mm256_set1_ps(-127.0f)), _mm256_set1_ps(127.0f))); +} + +// The quants match a scalar absmax / 127 per-block quant that rounds with lround (half away from zero), bit for bit. +// 1 / scale overflows for absmax below ~3.7e-37, so such a block quantizes to zeros. +// Dot lanes 2b and 2b + 1 multiply the first and the last 8 activations of block b, so corr is two plain sums. +CPU_MOE_TARGET("avx2,fma") +void quant_i8_avx2(const bf16_t* x, int K, int8_t* q, float* scale, int32_t* corr) { + const __m128i even_odd = _mm_setr_epi8(0, 2, 4, 6, 8, 10, 12, 14, 1, 3, 5, 7, 9, 11, 13, 15); + const __m256 absm = _mm256_castsi256_ps(_mm256_set1_epi32(0x7FFFFFFF)); const int nb = K / 16; - // Built-in default: min(4 KB, 2 rows) -- see nvfp4_pf_blocks(). An explicitly set - // env value is used verbatim so a per-machine sweep can reach operating points past - // the conservative default. - const int pfb = nvfp4_pf_blocks(); - const int pf = (pfb < 0) ? std::min(512, 2 * nb) : pfb; + for (int b = 0; b < nb; ++b) { + const int c0 = b & ~7, nbc = std::min(8, nb - c0); + int8_t* qe = q + (size_t)c0 * 16 + (size_t)(b - c0) * 8; + const __m256i raw = _mm256_loadu_si256(reinterpret_cast(x + (size_t)b * 16)); + const __m256 v0 = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(_mm256_castsi256_si128(raw)), 16)); + const __m256 v1 = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(_mm256_extracti128_si256(raw, 1)), 16)); + const __m256 m8 = _mm256_max_ps(_mm256_and_ps(v0, absm), _mm256_and_ps(v1, absm)); + __m128 m4 = _mm_max_ps(_mm256_castps256_ps128(m8), _mm256_extractf128_ps(m8, 1)); + m4 = _mm_max_ps(m4, _mm_movehl_ps(m4, m4)); + m4 = _mm_max_ss(m4, _mm_shuffle_ps(m4, m4, 0x55)); + const float amax = _mm_cvtss_f32(m4); + const float s = amax > 0.0f ? amax / 127.0f : 1.0f; + scale[b] = s; + const float inv_s = 1.0f / s; + const __m256 inv = _mm256_set1_ps(inv_s < INFINITY ? inv_s : 0.0f); + const __m256i i0 = w4a8_round256(_mm256_mul_ps(v0, inv)), i1 = w4a8_round256(_mm256_mul_ps(v1, inv)); + const __m256i p16 = _mm256_permute4x64_epi64(_mm256_packs_epi32(i0, i1), 0xD8); + const __m128i bytes = _mm_shuffle_epi8(_mm_packs_epi16(_mm256_castsi256_si128(p16), _mm256_extracti128_si256(p16, 1)), even_odd); + _mm_storel_epi64(reinterpret_cast<__m128i*>(qe), bytes); + _mm_storel_epi64(reinterpret_cast<__m128i*>(qe + nbc * 8), _mm_unpackhi_epi64(bytes, bytes)); + corr[2 * b] = -kE2M1Bias * hsum_epi32_256(i0); + corr[2 * b + 1] = -kE2M1Bias * hsum_epi32_256(i1); + } +} + +#ifdef CPU_MOE_HAS_AVX512VNNI + +// Weight scale * act scale of n <= 16 blocks, in 2^-8 units. +CPU_MOE_TARGET("avx512f,avx512bw,avx512vl") +inline __m512 w4a8_scales512(const uint8_t* scale, const float* asc, int n) { + const __mmask16 m = (__mmask16)((1u << n) - 1); + __m256i h = _mm256_cvtepi8_epi16(_mm_maskz_loadu_epi8(m, scale)); + h = _mm256_and_si256(_mm256_slli_epi16(h, 7), _mm256_set1_epi16(kE4m3F16Mask)); + return _mm512_mul_ps(_mm512_cvtph_ps(h), _mm512_maskz_loadu_ps(m, asc)); +} + +CPU_MOE_TARGET("avx512f,avx512bw,avx512vl") +inline void w4a8_split512(__m512i v, __m512i lut, __m512i& lo, __m512i& hi) { + const __m512i m4 = _mm512_set1_epi8(0x0F); + lo = _mm512_shuffle_epi8(lut, _mm512_and_si512(v, m4)); + hi = _mm512_shuffle_epi8(lut, _mm512_and_si512(_mm512_srli_epi16(v, 4), m4)); +} + +CPU_MOE_TARGET("avx512f,avx512bw,avx512vl,avx512vnni") +inline __m512i w4a8_dot512_vnni(__m512i corr, __m512i lo, __m512i qe, __m512i hi, __m512i qo) { + return _mm512_dpbusd_epi32(_mm512_dpbusd_epi32(corr, lo, qe), hi, qo); +} + +CPU_MOE_TARGET("avx512f,avx512bw,avx512vl,avx512vnni") +float dot_nvfp4_i8_avx512vnni(const uint8_t* packed, const uint8_t* scale, float global, const W4a8Act& a, int K) { + const __m512i lut = _mm512_broadcast_i32x4(_mm_load_si128(reinterpret_cast(kE2M1x2u))); + const __m512i lanes_lo = _mm512_set_epi32(7, 7, 6, 6, 5, 5, 4, 4, 3, 3, 2, 2, 1, 1, 0, 0); + const __m512i lanes_hi = _mm512_add_epi32(lanes_lo, _mm512_set1_epi32(8)); + const int nb = K / 16, pf = w4a8_pf(nb); + __m512 acc0 = _mm512_setzero_ps(), acc1 = _mm512_setzero_ps(); + __m512i lo, hi; int b = 0; - // Four independent 4-block groups per iter -> four accumulator chains hide the ~10-op - // decode->dpbusd->scale latency (1-2 chains leave the loop latency-bound: the per-core - // rate sat at ~60% of the core's achievable DRAM stream rate). for (; b + 16 <= nb; b += 16) { - // Prefetch the weight stream ahead: the interleaved decode lowers the L1-miss - // concurrency the HW prefetcher sustains on its own (134 -> 167 GB/s on Emerald Rapids). + const uint8_t* pk = packed + (size_t)b * 8; + // The HW prefetcher alone falls short here: with pf = 0 a GLM-5.3-Flash executor step is 1.27x slower on Emerald Rapids. if (pf > 0) { - _mm_prefetch(reinterpret_cast(packed + ((size_t)b + (size_t)pf) * 8), - _MM_HINT_T0); - _mm_prefetch(reinterpret_cast(packed + ((size_t)b + (size_t)pf) * 8 + 64), - _MM_HINT_T0); + _mm_prefetch(reinterpret_cast(pk + (size_t)pf * 8), _MM_HINT_T0); + _mm_prefetch(reinterpret_cast(pk + (size_t)pf * 8 + 64), _MM_HINT_T0); } - acc0 = _mm512_add_ps(acc0, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b, lut, idx, - mask0F, idxsc)); - acc1 = _mm512_add_ps(acc1, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b + 4, lut, idx, - mask0F, idxsc)); - acc2 = _mm512_add_ps(acc2, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b + 8, lut, idx, - mask0F, idxsc)); - acc3 = _mm512_add_ps(acc3, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b + 12, lut, idx, - mask0F, idxsc)); + const int8_t* q = a.q + (size_t)b * 16; + const int32_t* c = a.corr + (size_t)b * 2; + const __m512 sc = w4a8_scales512(scale + b, a.scale + b, 16); + w4a8_split512(_mm512_loadu_si512(pk), lut, lo, hi); + __m512i d = w4a8_dot512_vnni(_mm512_loadu_si512(c), lo, _mm512_loadu_si512(q), hi, _mm512_loadu_si512(q + 64)); + acc0 = _mm512_fmadd_ps(_mm512_cvtepi32_ps(d), _mm512_permutexvar_ps(lanes_lo, sc), acc0); + w4a8_split512(_mm512_loadu_si512(pk + 64), lut, lo, hi); + d = w4a8_dot512_vnni(_mm512_loadu_si512(c + 16), lo, _mm512_loadu_si512(q + 128), hi, _mm512_loadu_si512(q + 192)); + acc1 = _mm512_fmadd_ps(_mm512_cvtepi32_ps(d), _mm512_permutexvar_ps(lanes_hi, sc), acc1); } - for (; b + 4 <= nb; b += 4) - acc0 = _mm512_add_ps(acc0, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b, lut, idx, - mask0F, idxsc)); - float s = _mm512_reduce_add_ps( - _mm512_add_ps(_mm512_add_ps(acc0, acc1), _mm512_add_ps(acc2, acc3))); - for (; b < nb; ++b) { // tail (<4 remaining 16-K blocks) - const uint8_t* pk = packed + (size_t)b * 8; - const int8_t* ae = asi8 + (size_t)b * 16; const int8_t* ao = ae + 8; - int isum = 0; - for (int j = 0; j < 8; ++j) - isum += (int)kE2M1x2[pk[j] & 0xF] * (int)ae[j] + (int)kE2M1x2[pk[j] >> 4] * (int)ao[j]; - s += (e4m3[scale[b]] * asb[b]) * (float)isum; + // The tail runs chunks of n <= 8 blocks with masked loads: masked-off activations are zero, so those lanes add nothing. + for (; b < nb; b += 8) { + const int n = std::min(8, nb - b); + const __mmask64 mb = n == 8 ? ~0ULL : (1ULL << (8 * n)) - 1; + const int8_t* q = a.q + (size_t)b * 16; + const __m512 sc = w4a8_scales512(scale + b, a.scale + b, n); + w4a8_split512(_mm512_maskz_loadu_epi8(mb, packed + (size_t)b * 8), lut, lo, hi); + const __m512i c = _mm512_maskz_loadu_epi32((__mmask16)((1u << (2 * n)) - 1), a.corr + (size_t)b * 2); + const __m512i d = w4a8_dot512_vnni(c, lo, _mm512_maskz_loadu_epi8(mb, q), hi, _mm512_maskz_loadu_epi8(mb, q + 8 * n)); + acc0 = _mm512_fmadd_ps(_mm512_cvtepi32_ps(d), _mm512_permutexvar_ps(lanes_lo, sc), acc0); + } + return _mm512_reduce_add_ps(_mm512_add_ps(acc0, acc1)) * (0.5f * kE4m3ViaF16 * global); +} + +CPU_MOE_TARGET("avx512f,avx512bw,avx512vl") +void quant_i8_avx512(const bf16_t* x, int K, int8_t* q, float* scale, int32_t* corr) { + const __m128i even_odd = _mm_setr_epi8(0, 2, 4, 6, 8, 10, 12, 14, 1, 3, 5, 7, 9, 11, 13, 15); + const __m512 half = _mm512_set1_ps(0.5f); + const __m512i one = _mm512_castps_si512(_mm512_set1_ps(1.0f)), sign = _mm512_set1_epi32(INT32_MIN); + const __m512 lim = _mm512_set1_ps(127.0f), nlim = _mm512_set1_ps(-127.0f); + const int nb = K / 16; + for (int b = 0; b < nb; ++b) { + const int c0 = b & ~7, nbc = std::min(8, nb - c0); + int8_t* qe = q + (size_t)c0 * 16 + (size_t)(b - c0) * 8; + const __m256i raw = _mm256_loadu_si256(reinterpret_cast(x + (size_t)b * 16)); + const __m512 v = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(raw), 16)); + const float amax = _mm512_reduce_max_ps(_mm512_abs_ps(v)); + const float s = amax > 0.0f ? amax / 127.0f : 1.0f; + scale[b] = s; + const float inv = 1.0f / s; + const __m512 y = _mm512_mul_ps(v, _mm512_set1_ps(inv < INFINITY ? inv : 0.0f)); + const __m512 t = _mm512_roundscale_ps(y, _MM_FROUND_TO_ZERO | _MM_FROUND_NO_EXC); + const __mmask16 away = _mm512_cmp_ps_mask(_mm512_abs_ps(_mm512_sub_ps(y, t)), half, _CMP_GE_OQ); + const __m512 step = _mm512_castsi512_ps(_mm512_or_si512(_mm512_and_si512(_mm512_castps_si512(y), sign), one)); + const __m512 r = _mm512_mask_add_ps(t, away, t, step); + const __m512i ri = _mm512_cvtps_epi32(_mm512_min_ps(_mm512_max_ps(r, nlim), lim)); + const __m128i bytes = _mm_shuffle_epi8(_mm512_cvtepi32_epi8(ri), even_odd); + _mm_storel_epi64(reinterpret_cast<__m128i*>(qe), bytes); + _mm_storel_epi64(reinterpret_cast<__m128i*>(qe + nbc * 8), _mm_unpackhi_epi64(bytes, bytes)); + corr[2 * b] = -kE2M1Bias * hsum_epi32_256(_mm512_castsi512_si256(ri)); + corr[2 * b + 1] = -kE2M1Bias * hsum_epi32_256(_mm512_extracti64x4_epi64(ri, 1)); } - return s * (0.5f * global); } #endif // avx512vnni available #endif @@ -317,45 +400,48 @@ nvdot_fn select_nvdot() { return dot_nvfp4_scalar; } -// Best W4A8 (int8-activation) nvfp4 dot, or nullptr if no SIMD VNNI (caller keeps the -// faithful fp32 nvdot path). The scalar i8 dot exists only as a correctness reference. -nvi8dot_fn select_nvi8dot() { +struct W4a8Kernel { + nvi8dot_fn dot; + nvi8quant_fn quant; + const char* isa; +}; + +// Best W4A8 (int8-activation) nvfp4 dot for the picked ISA tier, or a null dot if that tier has no VNNI (caller keeps the faithful fp32 nvdot path). +W4a8Kernel select_nvi8dot() { #if CPU_MOE_X86 + const IsaTier t = pick_isa(); #if defined(CPU_MOE_HAS_AVX512VNNI) - if (cpu_has_avx512vnni()) return dot_nvfp4_i8_avx512vnni; + if (t >= ISA_AVX512 && cpu_has_avx512vnni()) return {dot_nvfp4_i8_avx512vnni, quant_i8_avx512, "+avx512vnni(nvfp4-w4a8)"}; #endif - if (cpu_has_avxvnni()) return dot_nvfp4_i8_vnni; + if (t >= ISA_AVX2 && cpu_has_avxvnni()) return {dot_nvfp4_i8_vnni, quant_i8_avx2, "+vnni(nvfp4-w4a8)"}; #endif - return nullptr; + return {}; } -// Quantize the pre-deinterleaved fp32 even/odd activations to per-16-block int8 in the -// [even(8),odd(8)] layout the VNNI dot expects. Done once per token/route (amortized over -// every output row), so a scalar pass is fine relative to the GEMV. -void quant_i8_pg16(const float* xe, const float* xo, int K, int8_t* asi8, float* asb) { - const int nb = K / 16; - for (int b = 0; b < nb; ++b) { - const float* xeb = xe + (size_t)b * 8; - const float* xob = xo + (size_t)b * 8; - float amax = 0.0f; - for (int j = 0; j < 8; ++j) - amax = std::max(amax, std::max(std::fabs(xeb[j]), std::fabs(xob[j]))); - const float s = amax > 0.0f ? amax / 127.0f : 1.0f; - asb[b] = s; - const float inv = 1.0f / s; - int8_t* ae = asi8 + (size_t)b * 16; - for (int j = 0; j < 8; ++j) { - int qe = (int)std::lround(xeb[j] * inv), qo = (int)std::lround(xob[j] * inv); - ae[j] = (int8_t)std::max(-127, std::min(127, qe)); - ae[8 + j] = (int8_t)std::max(-127, std::min(127, qo)); - } +struct W4a8Rows { + AlignedVector q; + AlignedVector scale; + AlignedVector corr; + int K = 0; + + void resize(size_t rows, int k) { + K = k; + q.resize(rows * k); + scale.resize(rows * (k / 16)); + corr.resize(rows * (k / 8)); } -} + void quant(nvi8quant_fn fn, size_t row, const bf16_t* x) { + fn(x, K, q.data() + row * K, scale.data() + row * (K / 16), corr.data() + row * (K / 8)); + } + W4a8Act at(size_t row) const { + return {q.data() + row * K, scale.data() + row * (K / 16), corr.data() + row * (K / 8)}; + } +}; class Nvfp4Format final : public MoeFormat { public: Nvfp4Format(const MoeParams& p, const FormatArgs& a) - : MoeFormat(p), nvdot_(select_nvdot()), nvi8dot_(select_nvi8dot()) { + : MoeFormat(p), nvdot_(select_nvdot()), w4a8_(select_nvi8dot()) { gate_up_ = a.bank("gate_up"); gu_scale_ = a.bank("gate_up_scale"); gu_global_ = a.bank("gate_up_global"); @@ -363,18 +449,17 @@ class Nvfp4Format final : public MoeFormat { dn_scale_ = a.bank("down_scale"); dn_global_ = a.bank("down_global"); const size_t tokens = p.max_tokens, routes = (size_t)p.max_tokens * p.top_k; - xe_.resize(tokens * (p.H / 2)); - xo_.resize(tokens * (p.H / 2)); - ge_.resize(routes * (p.I / 2)); - go_.resize(routes * (p.I / 2)); - if (nvi8dot_ != nullptr) { - xi8_.resize(tokens * p.H); - xas_.resize(tokens * (p.H / 16)); - gi8_.resize(routes * p.I); - gas_.resize(routes * (p.I / 16)); + if (w4a8_.dot != nullptr) { + xq_.resize(tokens, p.H); + gq_.resize(routes, p.I); + } else { + xe_.resize(tokens * (p.H / 2)); + xo_.resize(tokens * (p.H / 2)); + ge_.resize(routes * (p.I / 2)); + go_.resize(routes * (p.I / 2)); } isa_ = tier_name(pick_isa()); - if (nvi8dot_ != nullptr) isa_ += cpu_has_avx512vnni() ? "+avx512vnni(nvfp4-w4a8)" : "+vnni(nvfp4-w4a8)"; + if (w4a8_.dot != nullptr) isa_ += w4a8_.isa; for (int i = 0; i < 16; ++i) e2m1_lut_[i] = kE2M1[i]; for (int i = 0; i < 256; ++i) e4m3_lut_[i] = e4m3_decode((uint8_t)i); } @@ -383,21 +468,25 @@ class Nvfp4Format final : public MoeFormat { void prep_input(int tok, const bf16_t* x) override { + if (w4a8_.dot != nullptr) { + xq_.quant(w4a8_.quant, tok, x); + return; + } float* xe = xe_.data() + (size_t)tok * (p.H / 2); float* xo = xo_.data() + (size_t)tok * (p.H / 2); deinterleave_bf16_f32(x, xe, xo, p.H); - if (nvi8dot_ != nullptr) - quant_i8_pg16(xe, xo, p.H, xi8_.data() + (size_t)tok * p.H, xas_.data() + (size_t)tok * (p.H / 16)); } bool preps_inter() const override { return true; } void prep_inter(int64_t r, bf16_t* g) override { + if (w4a8_.dot != nullptr) { + gq_.quant(w4a8_.quant, r, g); + return; + } float* ge = ge_.data() + (size_t)r * (p.I / 2); float* go = go_.data() + (size_t)r * (p.I / 2); deinterleave_bf16_f32(g, ge, go, p.I); - if (nvi8dot_ != nullptr) - quant_i8_pg16(ge, go, p.I, gi8_.data() + (size_t)r * p.I, gas_.data() + (size_t)r * (p.I / 16)); } void gate_up(int layer, int e, int tok, int i0, int i1, const bf16_t* /*x*/, float w_in, @@ -406,15 +495,18 @@ class Nvfp4Format final : public MoeFormat { const uint8_t* packed = static_cast(tbl_at(gate_up_, layer)); const uint8_t* scale = static_cast(tbl_at(gu_scale_, layer)); const float* global = static_cast(tbl_at(gu_global_, layer)); + if (w4a8_.dot != nullptr) { + const W4a8Act act = xq_.at(tok); + gated_rows(p, i0, i1, w_in, g, [&](int row) { + const size_t r = (size_t)e * (2 * p.I) + row; + return w4a8_.dot(packed + r * (size_t)(H / 2), scale + r * (size_t)(H / 16), global[r], act, H); + }); + return; + } const float* xe = xe_.data() + (size_t)tok * (H / 2); const float* xo = xo_.data() + (size_t)tok * (H / 2); - const int8_t* xi8 = nvi8dot_ != nullptr ? xi8_.data() + (size_t)tok * H : nullptr; - const float* xas = nvi8dot_ != nullptr ? xas_.data() + (size_t)tok * (H / 16) : nullptr; gated_rows(p, i0, i1, w_in, g, [&](int row) { const size_t r = (size_t)e * (2 * p.I) + row; - if (nvi8dot_ != nullptr) - return nvi8dot_(packed + r * (size_t)(H / 2), scale + r * (size_t)(H / 16), - global[r], xi8, H, e4m3_lut_, xas); return nvdot_(packed + r * (size_t)(H / 2), scale + r * (size_t)(H / 16), global[r], xe, xo, H, e2m1_lut_, e4m3_lut_); }); @@ -429,9 +521,8 @@ class Nvfp4Format final : public MoeFormat { for (int j = 0; j < n; ++j) { const size_t r = (size_t)routes[j].r; const size_t row = (size_t)routes[j].e * p.H + h; - out[j * kRowBlock + h - h0] = nvi8dot_ != nullptr - ? nvi8dot_(packed + row * (size_t)(I / 2), scale + row * (size_t)(I / 16), global[row], - gi8_.data() + r * I, I, e4m3_lut_, gas_.data() + r * (I / 16)) + out[j * kRowBlock + h - h0] = w4a8_.dot != nullptr + ? w4a8_.dot(packed + row * (size_t)(I / 2), scale + row * (size_t)(I / 16), global[row], gq_.at(r), I) : nvdot_(packed + row * (size_t)(I / 2), scale + row * (size_t)(I / 16), global[row], ge_.data() + r * (I / 2), go_.data() + r * (I / 2), I, e2m1_lut_, e4m3_lut_); } @@ -440,7 +531,7 @@ class Nvfp4Format final : public MoeFormat { private: nvdot_fn nvdot_; - nvi8dot_fn nvi8dot_; // AVX-VNNI W4A8 nvfp4 dot (nullptr -> use fp32 nvdot) + const W4a8Kernel w4a8_; // null dot -> use fp32 nvdot // ModelOpt rows: packed e2m1 [E,2I,H/2], e4m3 scales per 16 along K [E,2I,H/16], fp32 row globals [E,2I]; down likewise with K = I. const uint64_t* gate_up_ = nullptr; const uint64_t* gu_scale_ = nullptr; @@ -451,10 +542,9 @@ class Nvfp4Format final : public MoeFormat { std::string isa_; float e2m1_lut_[16]; float e4m3_lut_[256]; - // Activations pre-deinterleaved to fp32 even/odd for the row-major dot; with VNNI also per-16-block int8 + scale. + // fp32 even/odd activations for the fp32 dot, or the W4A8 int8 rows. AlignedVector xe_, xo_, ge_, go_; - AlignedVector xi8_, gi8_; - AlignedVector xas_, gas_; + W4a8Rows xq_, gq_; }; } // namespace diff --git a/python/freetoken/moe/benchbw.py b/python/freetoken/moe/benchbw.py index 0225cc4fa2..911ca664b3 100644 --- a/python/freetoken/moe/benchbw.py +++ b/python/freetoken/moe/benchbw.py @@ -699,8 +699,9 @@ def _bench_format(fmt: str, wl: Workload, device: torch.device, threshold: float # The tier label comes from the bf16 dot chooser, so it can over-report tiers # the actual per-format kernel doesn't distinguish. Disclose the real story. if fmt == "nvfp4": # nvfp4 is the only use_vnni format - _note(entry, "isa tier labels are nominal: this W4A8 kernel rides AVX-VNNI " - "regardless of the forced tier") + _note(entry, "isa tier labels are nominal: avx512bf16 == avx512f here " + "(avx512 runs W4A8 with AVX512-VNNI, avx2 with AVX-VNNI; " + "a tier without its VNNI runs fp32 W4A16)") elif fmt in ("mxfp4_triton", "ds_fp4", "fp8_block"): _note(entry, "isa tier labels are nominal: avx512bf16 == avx512f here " "(3 real kernels: scalar/avx2/avx512, no bf16/VNNI variant)") diff --git a/tests/moe/test_cpu_moe.py b/tests/moe/test_cpu_moe.py index edd9ddc9df..89f1daa32e 100644 --- a/tests/moe/test_cpu_moe.py +++ b/tests/moe/test_cpu_moe.py @@ -111,24 +111,28 @@ def rows(OUT, IN): ) +# FREETOKEN_CPU_MOE_ISA tier -> the tier the executor reports. A tier the CPU or the build lacks is capped to a lower one, which turns into a skip. +_NVFP4_TIERS = {"scalar": "scalar", "avx2": "avx2", "avx512": "avx512f"} + + +@pytest.mark.parametrize("isa", list(_NVFP4_TIERS)) @pytest.mark.parametrize("bs", [1, 3, 8]) -def test_cpu_decode_nvfp4_matches_dequant_then_gpu(bs): - """CPU inline-dequant NVFP4 GEMV vs. canonical dequant_nvfp4 + bf16 GPU decode. +def test_cpu_decode_nvfp4_matches_dequant_then_gpu(bs, isa, monkeypatch): + """CPU NVFP4 GEMV vs. canonical dequant_nvfp4 + bf16 GPU decode, per ISA tier. - Both paths read byte-identical quantized banks; the CPU dequantizes inside the - K-loop (fp32), the reference materializes bf16 weights first, so the only spread - is weight bf16-rounding + reduction order -> tight relative tolerance. + Both paths read byte-identical quantized banks. The avx2 / avx512 tiers run W4A8 when the CPU has their VNNI, and the int8 activation rounding dominates the spread; the fp32 dots differ only by weight bf16-rounding + reduction order. """ from freetoken.kernel.triton.nvfp4_dequant import dequant_nvfp4 from freetoken.moe.cpu_executor import CpuMoeExecutor from freetoken.moe.fused import fused_experts_decode_impl torch.manual_seed(100 + bs) - L, E, H, I, top_k = 3, 16, 1024, 512, 4 + L, E, H, I, top_k = 3, 16, 1024, 752, 4 # 47 blocks of 16: the W4A8 down dots run their partial-chunk and scalar tails layer = 1 dev = torch.device("cuda") cache = _make_nvfp4_cache(L, E, H, I) + monkeypatch.setenv("FREETOKEN_CPU_MOE_ISA", isa) ex = CpuMoeExecutor( cache, top_k=top_k, @@ -138,6 +142,8 @@ def test_cpu_decode_nvfp4_matches_dequant_then_gpu(bs): max_tokens=bs, device=dev, ) + if ex.isa.split("+")[0] != _NVFP4_TIERS[isa]: + pytest.skip(f"asked for {isa}, host/build selected {ex.isa}: not supported here") hidden = torch.randn(bs, H, device=dev, dtype=torch.bfloat16) ids = torch.stack([torch.randperm(E, device=dev)[:top_k] for _ in range(bs)]).to(torch.int32)