From b34647906550f6e9dc65c55ecfb36b697d5e1f43 Mon Sep 17 00:00:00 2001 From: Yuwen Zhou Date: Fri, 18 Sep 2026 17:59:42 +0800 Subject: [PATCH] [CPU] Add CPU FP8 W8A8 linear/MoE support (#49942) Signed-off-by: yuwenzho Co-authored-by: Li, Jiang --- cmake/cpu_extension.cmake | 5 + csrc/cpu/cpu_attn.cpp | 13 +- csrc/cpu/cpu_isa.cpp | 20 + csrc/cpu/sgl-kernels/common.h | 83 ++ csrc/cpu/sgl-kernels/gemm.h | 54 +- csrc/cpu/sgl-kernels/gemm_fp8_w8a8.cpp | 1093 +++++++++++++++++ csrc/cpu/sgl-kernels/moe.cpp | 100 +- csrc/cpu/sgl-kernels/moe_fp8_w8a8.cpp | 432 +++++++ csrc/cpu/torch_bindings.cpp | 37 +- tests/kernels/moe/test_cpu_quant_fused_moe.py | 198 +++ .../quantization/test_cpu_fp8_scaled_mm.py | 229 +++- vllm/_custom_ops.py | 44 + .../model_executor/kernels/linear/__init__.py | 4 + .../kernels/linear/scaled_mm/__init__.py | 2 + .../kernels/linear/scaled_mm/cpu.py | 310 ++++- .../layers/attention/mla_attention.py | 4 +- .../layers/fused_moe/experts/cpu_moe.py | 217 ++++ .../layers/fused_moe/oracle/fp8.py | 32 +- .../layers/quantization/online/fp8.py | 6 +- vllm/v1/attention/backends/mla/amx_mla.py | 87 +- 20 files changed, 2924 insertions(+), 46 deletions(-) create mode 100644 csrc/cpu/cpu_isa.cpp create mode 100644 csrc/cpu/sgl-kernels/gemm_fp8_w8a8.cpp create mode 100644 csrc/cpu/sgl-kernels/moe_fp8_w8a8.cpp diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index 6359a96887..cca2a98cee 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -488,6 +488,7 @@ set(VLLM_EXT_SRC "csrc/moe/dynamic_4bit_int_moe_cpu.cpp" "csrc/cpu/cpu_fused_moe.cpp" "csrc/cpu/cpu_attn.cpp" + "csrc/cpu/cpu_isa.cpp" "csrc/cpu/torch_bindings.cpp") if (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64" AND VLLM_RVV_VLEN AND @@ -546,11 +547,13 @@ if (ENABLE_X86_ISA) "csrc/cpu/sgl-kernels/gemm.cpp" "csrc/cpu/sgl-kernels/gemm_int8.cpp" "csrc/cpu/sgl-kernels/gemm_fp8.cpp" + "csrc/cpu/sgl-kernels/gemm_fp8_w8a8.cpp" "csrc/cpu/sgl-kernels/gemm_int4.cpp" "csrc/cpu/sgl-kernels/moe.cpp" "csrc/cpu/sgl-kernels/moe_int8.cpp" "csrc/cpu/sgl-kernels/moe_int4.cpp" "csrc/cpu/sgl-kernels/moe_fp8.cpp" + "csrc/cpu/sgl-kernels/moe_fp8_w8a8.cpp" "csrc/cpu/sgl-kernels/bmm.cpp" "csrc/cpu/sgl-kernels/decode.cpp" "csrc/cpu/sgl-kernels/extend.cpp" @@ -571,6 +574,7 @@ if (ENABLE_X86_ISA) "csrc/cpu/utils.cpp" "csrc/cpu/spec_decode_utils.cpp" "csrc/cpu/cpu_attn.cpp" + "csrc/cpu/cpu_isa.cpp" "csrc/cpu/dnnl_kernels.cpp" "csrc/cpu/mamba_cpu.cpp" "csrc/cpu/torch_bindings.cpp" @@ -587,6 +591,7 @@ if (ENABLE_X86_ISA) "csrc/cpu/utils.cpp" "csrc/cpu/spec_decode_utils.cpp" "csrc/cpu/cpu_attn.cpp" + "csrc/cpu/cpu_isa.cpp" "csrc/cpu/mamba_cpu.cpp" "csrc/cpu/dnnl_kernels.cpp" "csrc/cpu/torch_bindings.cpp" diff --git a/csrc/cpu/cpu_attn.cpp b/csrc/cpu/cpu_attn.cpp index b31a6faa01..ea0b4e49b2 100644 --- a/csrc/cpu/cpu_attn.cpp +++ b/csrc/cpu/cpu_attn.cpp @@ -1,16 +1,7 @@ #include "cpu_attn_dispatch_generated.h" -#if defined(__x86_64__) || defined(_M_X64) - #include -// Runtime check for AMX-FP8 (Diamond Rapids): CPUID leaf 7, subleaf 1, EAX[21]. -static bool runtime_has_amx_fp8() { - unsigned int eax = 0, ebx = 0, ecx = 0, edx = 0; - if (__get_cpuid_count(7, 1, &eax, &ebx, &ecx, &edx)) { - return (eax >> 21) & 1u; - } - return false; -} -#endif // x86_64 +// Runtime check for AMX-FP8, implemented in cpu_isa.cpp. +extern bool runtime_has_amx_fp8(); // Maps kv_cache_dtype string to Fp8KVCacheDataType enum. // "auto" -> kAuto(0); "fp8"/"fp8_e4m3" -> kFp8E4M3; "fp8_e5m2" -> kFp8E5M2. diff --git a/csrc/cpu/cpu_isa.cpp b/csrc/cpu/cpu_isa.cpp new file mode 100644 index 0000000000..1c4a8faf19 --- /dev/null +++ b/csrc/cpu/cpu_isa.cpp @@ -0,0 +1,20 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +// AMX-FP8 is reported at CPUID leaf 0x1E, subleaf 1, EAX[4]. +#if defined(__x86_64__) || defined(_M_X64) + #include + +bool runtime_has_amx_fp8() { + unsigned int eax = 0, ebx = 0, ecx = 0, edx = 0; + if (__get_cpuid_max(0, nullptr) < 0x1e) return false; + if (__get_cpuid_count(0x1e, 1, &eax, &ebx, &ecx, &edx)) { + return (eax >> 4) & 1u; + } + return false; +} + +bool cpu_has_amx_fp8() { return runtime_has_amx_fp8(); } +#else +bool cpu_has_amx_fp8() { return false; } +#endif diff --git a/csrc/cpu/sgl-kernels/common.h b/csrc/cpu/sgl-kernels/common.h index 6f27ba2d5c..de1260bb33 100644 --- a/csrc/cpu/sgl-kernels/common.h +++ b/csrc/cpu/sgl-kernels/common.h @@ -12,11 +12,45 @@ #include #include #include +#include #if defined(_OPENMP) #include #endif +// Adapted from sglang: dispatch utilities for FP8 W8A8 kernels +#define AT_DISPATCH_OUT_TYPES(TYPE, NAME, ...) \ + AT_DISPATCH_SWITCH( \ + TYPE, \ + NAME, \ + AT_PRIVATE_CASE_TYPE_USING_HINT(at::ScalarType::Float, out_t, __VA_ARGS__) \ + AT_PRIVATE_CASE_TYPE_USING_HINT(at::ScalarType::BFloat16, out_t, __VA_ARGS__) \ + AT_PRIVATE_CASE_TYPE_USING_HINT(at::ScalarType::Half, out_t, __VA_ARGS__)) + +#define AT_DISPATCH_CASE_ENUM_NO_RETURN(VALUE, TYPE, HINT, ...) \ + case VALUE: { \ + constexpr TYPE HINT = VALUE; \ + __VA_ARGS__; \ + break; \ + } + +#define AT_DISPATCH_BOOL_NO_RETURN(VALUE, NAME, HINT, ...) \ + [&]() { \ + switch (VALUE) { \ + AT_DISPATCH_CASE_ENUM_NO_RETURN(true, bool, HINT, __VA_ARGS__) \ + AT_DISPATCH_CASE_ENUM_NO_RETURN(false, bool, HINT, __VA_ARGS__) \ + } \ + }() + +#define AT_DISPATCH_QUANT_MODE_NO_RETURN(MODE, NAME, HINT, ...) \ + [&]() { \ + switch (MODE) { \ + AT_DISPATCH_CASE_ENUM_NO_RETURN(PER_TENSOR, int, HINT, __VA_ARGS__) \ + AT_DISPATCH_CASE_ENUM_NO_RETURN(PER_ROW, int, HINT, __VA_ARGS__) \ + AT_DISPATCH_CASE_ENUM_NO_RETURN(PER_GROUP, int, HINT, __VA_ARGS__) \ + } \ + }() + namespace { // dispatch bool @@ -472,4 +506,53 @@ inline T* conditional_data_ptr(const std::optional& opt) { return opt.has_value() ? opt.value().data_ptr() : nullptr; } +// Adapted from sglang: M blocking utilities for FP8 W8A8 kernels +inline int64_t get_m_block(int64_t M) { + if (M <= 48) { + return M; + } else if (M < 64) { + return 32; + } else if (M < 96) { + return 64; + } else { + return 128; + } +} + +inline std::tuple get_m_blocking(int64_t M) { + bool parallel_on_M = M > 128; + int64_t block_m = get_m_block(M); + int64_t Mc = (M + block_m - 1) / block_m; + int64_t Mc_parallel = parallel_on_M ? Mc : 1; + return std::make_tuple(parallel_on_M, block_m, Mc, Mc_parallel); +} + +#if defined(CPU_CAPABILITY_AVX512) +template +void zero_buffer(T* data, int64_t size) { + const int32_t vec_size = at::vec::Vectorized::size(); + auto zero_vec = at::vec::Vectorized(0); + int64_t d = 0; + for (; d < size - (size % vec_size); d += vec_size) { + zero_vec.store(data + d); + } + if (d < size) { + zero_vec.store(data + d, size - d); + } +} +#else +template +void zero_buffer(T* data, int64_t size) { + memset(data, 0, sizeof(T) * size); +} +#endif + +#if defined(__x86_64__) +bool avx10_2_available() { + // __builtin_cpu_supports returns the masked feature bit rather than 0/1. + static const bool available = __builtin_cpu_supports("avx10.2"); + return available; +} +#endif + } // anonymous namespace diff --git a/csrc/cpu/sgl-kernels/gemm.h b/csrc/cpu/sgl-kernels/gemm.h index db7ad884d0..dc5f9c585d 100644 --- a/csrc/cpu/sgl-kernels/gemm.h +++ b/csrc/cpu/sgl-kernels/gemm.h @@ -94,7 +94,7 @@ enum class CPUActMethod : int { gelu_and_mul = 2, }; -enum class CPUQuantMethod : int64_t { BF16 = 0, INT8_W8A8 = 1, FP8_W8A16 = 2, INT4_W4A8 = 3, MXFP4 = 4 }; +enum class CPUQuantMethod : int64_t { BF16 = 0, INT8_W8A8 = 1, FP8_W8A16 = 2, INT4_W4A8 = 3, MXFP4 = 4, FP8_W8A8 = 5 }; constexpr bool operator==(CPUQuantMethod a, int64_t b) { return static_cast(a) == b; @@ -395,3 +395,55 @@ void tinygemm_kernel( bool brg, int64_t block_size_K, bool do_unpack = true); + +// FP8 W8A8 tinygemm kernel (defined in gemm_fp8_w8a8.cpp) +void tinygemm_kernel( + float* C, + const at::Float8_e4m3fn* A, + const float* scales_a, + const at::Float8_e4m3fn* B, + const float* scales_b, + int64_t M, + int64_t K, + int64_t lda, + int64_t ldc, + int64_t ldsa, + float* ukernel_buf, + at::BFloat16* dqA_buf, + at::BFloat16* dqB_buf); + +// FP8 W8A8 quantize activation (defined in gemm_fp8_w8a8.cpp) +std::tuple quantize_fp8e4m3_vec( + const at::Tensor& t, + bool channelwise, + c10::optional scale_opt); + +// FP8 W8A8 fused MoE kernel (defined in moe_fp8_w8a8.cpp) +template +void fused_experts_fp8_a8_kernel_impl( + scalar_t* __restrict__ output, + scalar_t* __restrict__ ic0, + scalar_t* __restrict__ ic1, + scalar_t* __restrict__ ic2, + at::Float8_e4m3fn* __restrict__ A_tmp, + scalar_t* __restrict__ B_tmp, + float* __restrict__ C_tmp, + float* __restrict__ Ukernel_tmp, + const at::Float8_e4m3fn* __restrict__ input, + const at::Float8_e4m3fn* __restrict__ packed_w1, + const at::Float8_e4m3fn* __restrict__ packed_w2, + const float* __restrict__ As, + const float* __restrict__ w1s, + const float* __restrict__ w2s, + int64_t block_size_N, + int64_t block_size_K, + const float* __restrict__ topk_weights, + const int32_t* __restrict__ sorted_ids, + const int32_t* __restrict__ expert_ids, + const int32_t* __restrict__ offsets, + int64_t M, + int64_t N, + int64_t K, + int64_t E, + int64_t topk, + int64_t num_tokens_post_pad); diff --git a/csrc/cpu/sgl-kernels/gemm_fp8_w8a8.cpp b/csrc/cpu/sgl-kernels/gemm_fp8_w8a8.cpp new file mode 100644 index 0000000000..854566d682 --- /dev/null +++ b/csrc/cpu/sgl-kernels/gemm_fp8_w8a8.cpp @@ -0,0 +1,1093 @@ +// Adapted from +// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu + +// clang-format off + +#include +#include +#include + +#include "common.h" +#include "vec.h" + +// Runtime check for AMX-FP8, implemented in cpu_isa.cpp. +extern bool cpu_has_amx_fp8(); + +namespace { + +#define BLOCK_N 32 + +#define PER_TENSOR 1 +#define PER_ROW 2 +#define PER_GROUP 3 + +static bool cpublas_checked = false; +static bool cpublas_can_pack = false; + +bool cpublas_could_pack() { + // the could_pack check requires AMX support implicitly + if (cpublas_checked) { + return cpublas_can_pack; + } + cpublas_can_pack = cpu_has_amx_fp8() && at::native::cpublas::could_pack(at::kFloat8_e4m3fn); + cpublas_checked = true; + return cpublas_can_pack; +} + +#if defined(CPU_CAPABILITY_AVX512) +static void cvt_f8e4m3_to_bf16( + const at::Float8_e4m3fn* __restrict__ in, at::BFloat16* out, int64_t rows, int64_t cols, int64_t stride) { + if (stride == cols) { + // A contiguous buffer + size_t len = rows * cols; + size_t i = 0; + for (; i < len; i += 32) { + __m256i fp8_vec = _mm256_loadu_si256((__m256i*)&in[i]); + __m512bh bf16_vec = cvt_e4m3_bf16_intrinsic_no_nan(fp8_vec); + _mm512_storeu_si512((__m512i*)(out + i), (__m512i)bf16_vec); + } + for (; i < len; ++i) { + out[i] = (at::BFloat16)in[i]; + } + } else { + // Non-contiguous. Access each row with stride + TORCH_CHECK(stride > cols); + for (int r = 0; r < rows; ++r) { + size_t i = 0; + size_t vec_len = cols / 32 * 32; + for (; i < vec_len; i += 32) { + __m256i fp8_vec = _mm256_loadu_si256((__m256i*)&in[r * stride + i]); + __m512bh bf16_vec = cvt_e4m3_bf16_intrinsic_no_nan(fp8_vec); + _mm512_storeu_si512((__m512i*)(out + r * cols + i), (__m512i)bf16_vec); + } + for (; i < cols; ++i) { + out[r * cols + i] = (at::BFloat16)in[r * stride + i]; + } + } + } +} + +// accumulate and store result to buffer +// if act/wei are per_group quantized, apply scales +template +static void _accumulate_result( + float* __restrict__ output, + const float* __restrict__ input, + const float* __restrict__ scale_a, + const float* __restrict__ scale_b, + int M, + int ldi, + int ldo, + int ldsa = 1) { + float a_scale, b_scale; + __m512 va_scale; + __m512 vb_scale; + for (int m = 0; m < M; ++m) { + if constexpr (act_quant_mode == PER_GROUP) { + a_scale = *(scale_a + m * ldsa); + va_scale = _mm512_set1_ps(a_scale); + } + constexpr int N_UNROLL = N / 16; + c10::ForcedUnroll{}([&](auto i) { + constexpr int n = i * 16; + __m512 vc_f = _mm512_loadu_ps(input + m * ldi + n); + if constexpr (act_quant_mode == PER_GROUP) { + vc_f = _mm512_mul_ps(vc_f, va_scale); + } + if constexpr (wei_quant_mode == PER_GROUP) { + vb_scale = _mm512_loadu_ps(scale_b + n); + vc_f = _mm512_mul_ps(vc_f, vb_scale); + } + if constexpr (accum) { + __m512 vo = _mm512_loadu_ps(output + m * ldo + n); + _mm512_storeu_ps(output + m * ldo + n, _mm512_add_ps(vo, vc_f)); + } else { + _mm512_storeu_ps(output + m * ldo + n, vc_f); + } + }); + constexpr int tail_start = N / 16 * 16; + for (int n = tail_start; n < N; ++n) { + float dq_val = input[m * ldi + n]; + if constexpr (act_quant_mode == PER_GROUP) { + dq_val = dq_val * a_scale; + } + if constexpr (wei_quant_mode == PER_GROUP) { + b_scale = scale_b[n]; + dq_val = dq_val * b_scale; + } + if constexpr (accum) { + output[m * ldo + n] += dq_val; + } else { + output[m * ldo + n] = dq_val; + } + } + } +} + +// Store result to output buffer with dtype conversion +// If act/wei are per_row or per_tensor quantized, apply scales +// If bias is not null, add bias +template +inline void store_out( + const float* y_buf, + out_dtype* c_ptr, + int64_t M, + int64_t lda, + const float* scales_a, + const float* scales_b, + const float* bias) { + float a_scale = 1.0, b_scale = 1.0; + __m512 va_scale, vb_scale; + if constexpr (act_quant_mode == PER_TENSOR) { + a_scale = *scales_a; + } + if constexpr (wei_quant_mode == PER_TENSOR) { + b_scale = *scales_b; + vb_scale = _mm512_set1_ps(b_scale); + } + for (int i = 0; i < M; ++i) { + if constexpr (act_quant_mode == PER_ROW) { + a_scale = *(scales_a + i); + } + if constexpr (act_quant_mode != PER_GROUP) { + va_scale = _mm512_set1_ps(a_scale); + } + constexpr int N_UNROLL = N / 16; + c10::ForcedUnroll{}([&](auto idx) { + constexpr int j = idx * 16; + __m512 y_vec = _mm512_loadu_ps(y_buf + i * N + j); + __m512 bias_vec = bias ? _mm512_loadu_ps(bias + j) : _mm512_setzero_ps(); + if constexpr (act_quant_mode != PER_GROUP) { + y_vec = _mm512_mul_ps(y_vec, va_scale); + } + if constexpr (wei_quant_mode == PER_ROW) { + vb_scale = _mm512_loadu_ps(scales_b + j); + } + if constexpr (wei_quant_mode != PER_GROUP) { + y_vec = _mm512_mul_ps(y_vec, vb_scale); + } + y_vec = _mm512_add_ps(y_vec, bias_vec); + if constexpr (std::is_same::value) { + _mm512_storeu_ps(c_ptr + i * lda + j, y_vec); + } else if constexpr (std::is_same::value) { + __m256i y_bf16_vec = at::vec::cvtfp32_bf16(y_vec); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c_ptr + i * lda + j), y_bf16_vec); + } else if constexpr (std::is_same::value) { + __m256i y_fp16_vec = at::vec::cvtfp32_fp16(y_vec); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c_ptr + i * lda + j), y_fp16_vec); + } else { + TORCH_CHECK(false, "Unsupported output dtype"); + } + }); + constexpr int tail_start = N / 16 * 16; + for (int j = tail_start; j < N; ++j) { + if constexpr (wei_quant_mode == PER_ROW) { + b_scale = scales_b[j]; + } + c_ptr[i * lda + j] = static_cast(y_buf[i * N + j] * a_scale * b_scale); + } + } // for M +} + +#else // no AVX512 + +static void cvt_f8e4m3_to_bf16( + const at::Float8_e4m3fn* __restrict__ in, at::BFloat16* out, int64_t rows, int64_t cols, int64_t stride) { + for (int r = 0; r < rows; ++r) { + for (int c = 0; c < cols; ++c) { + out[r * cols + c] = (at::BFloat16)in[r * stride + c]; + } + } +} + +// Store result to output buffer with dtype conversion +// If act/wei are per_row or per_tensor quantized, apply scales +// If bias is not null, add bias +template +inline void store_out( + const float* y_buf, + out_dtype* c_ptr, + int64_t M, + int64_t lda, + const float* scales_a, + const float* scales_b, + const float* bias) { + float a_scale = 1.0, b_scale = 1.0; + if constexpr (act_quant_mode == PER_TENSOR) { + a_scale = *scales_a; + } + if constexpr (wei_quant_mode == PER_TENSOR) { + b_scale = *scales_b; + } + for (int i = 0; i < M; ++i) { + if constexpr (act_quant_mode == PER_ROW) { + a_scale = *(scales_a + i); + } + for (int j = 0; j < N; ++j) { + if constexpr (wei_quant_mode == PER_ROW) { + b_scale = scales_b[j]; + } + c_ptr[i * lda + j] = static_cast(y_buf[i * N + j] * a_scale * b_scale); + } + } // for M +} + +#endif // CPU_CAPABILITY_AVX512 + +template +void _micro_gemm( + float* C, + const at::Float8_e4m3fn* A, + const float* scales_a, + const at::Float8_e4m3fn* B, + const float* scales_b, + int64_t M, + int64_t K, + int64_t lda, + int64_t ldc, + int64_t ldsa, + float* ukernel_buf, + at::BFloat16* dqA_buf, + at::BFloat16* dqB_buf) { + // If FP8 brgemm is not available, convert A/B to bf16 for computation + // Compute GEMM fp8 * fp8 -> fp32 (or bf16 * bf16 -> fp32) + // If per_group quant, apply scales. Otherwise, don't apply scales here + // Finally accumulate and store results +#if defined(CPU_CAPABILITY_AVX512) + if constexpr (cpublas_can_pack) { +#ifdef CPUBLAS_BRGEMM_F8F8F32 + at::native::cpublas::brgemm( + M, N, K, lda /*lda*/, N /*ldb*/, N /*ldc*/, false /* add_C */, A, B, ukernel_buf, true /* is_vnni */); +#else + cvt_f8e4m3_to_bf16(A, dqA_buf, M, K, lda); + cvt_f8e4m3_to_bf16(B, dqB_buf, K, N, N); + at::native::cpublas::brgemm( + M, N, K, K /*lda*/, N /*ldb*/, N /*ldc*/, false /* add_C */, dqA_buf, dqB_buf, ukernel_buf, true /* is_vnni */); +#endif + _mm_prefetch(B + N * (K + 128), _MM_HINT_T0); + _mm_prefetch(A + K + 128, _MM_HINT_T0); + _accumulate_result( + C, ukernel_buf, scales_a, scales_b, M, N /*ldi*/, ldc, ldsa); + } else +#endif + { + for (int64_t i = 0; i < M; ++i) { + for (int64_t j = 0; j < N; ++j) { + float sum = 0; + for (int64_t k = 0; k < K; ++k) { + sum += ((float)A[i * lda + k] * (float)B[k * N + j]); + } + if constexpr (act_quant_mode == PER_GROUP) { + sum *= scales_a[i * ldsa]; + } + if constexpr (wei_quant_mode == PER_GROUP) { + sum *= scales_b[j]; + } + C[i * ldc + j] += sum; + } + } + } +} + +template +void _float8_linear_impl( + const at::Tensor& input, + const at::Tensor& input_scales, + const at::Tensor& weight, + const at::Tensor& weight_scales, + const std::optional& bias, + at::Tensor& output) { + // input shape = [..., K] + // input is per token quantized + int64_t K = input.size(-1); + auto input_view = input.view({-1, K}); + int64_t M = input_view.size(0); + + // weight shape = [Nc, Kc, block_k, block_n] + // scales shape = [Nc, G, block_n] + int64_t Nc = weight.size(0); + int64_t Kc = wei_quant_mode != PER_GROUP ? 1 : weight.size(1); + int64_t block_k = wei_quant_mode != PER_GROUP ? weight.size(1) * weight.size(2) : weight.size(2); + constexpr int64_t block_n = BLOCK_N; + TORCH_CHECK(weight.size(3) == block_n, "Float8 linear: unexpected weight shape"); + int64_t N = Nc * block_n; + TORCH_CHECK(K == Kc * block_k, "Float8 linear: weight and input shapes mismatch"); + auto [parallel_on_M, block_m, Mc, Mc_parallel] = get_m_blocking(M); + + // scales shape = [Nc, G, block_n] + int64_t num_groups = wei_quant_mode == PER_TENSOR ? 1 : weight_scales.size(1); + TORCH_CHECK(K % num_groups == 0, "K should be divisible by num_groups"); + int64_t group_size = K / num_groups; + TORCH_CHECK(group_size % block_k == 0, "Float8 linear: group_size should be divisible by block_k"); + int64_t block_per_group = group_size / block_k; + TORCH_CHECK( + input_scales.numel() == 1 || input_scales.numel() == M || input_scales.numel() == M * num_groups, + "Float8 linear: unexpected input scales shape"); + auto ldsa = act_quant_mode == PER_TENSOR ? 0 : act_quant_mode == PER_ROW ? 1 : num_groups; + + const at::Float8_e4m3fn* a_ptr = input_view.data_ptr(); + const float* a_scales_ptr = input_scales.data_ptr(); + const at::Float8_e4m3fn* b_ptr = weight.data_ptr(); + const float* b_scales_ptr = weight_scales.data_ptr(); + out_dtype* c_ptr = output.data_ptr(); + const float* bias_ptr = bias.has_value() ? bias.value().data_ptr() : nullptr; + + int64_t block_size = block_m * block_n; + int64_t num_thread = at::get_num_threads(); + at::Tensor y_buffer = at::empty({num_thread, block_size}, output.options().dtype(at::kFloat)); + // Create buffer for brgemm output and dqA/dqB (optional) +#if defined(CPU_CAPABILITY_AVX512) + // buffer for brgemm output in float32 + int64_t buffer_size = block_size * 2; // float32 = bfloat16 * 2 +#ifndef CPUBLAS_BRGEMM_F8F8F32 + // buffers for dqA & dqB in bf16 + buffer_size += (block_k * block_n + block_m * block_k); +#endif + at::Tensor micro_gemm_buffer = at::empty({num_thread, buffer_size}, output.options().dtype(at::kBFloat16)); +#endif + + parallel_2d(Mc, Nc, [&](int64_t mc0, int64_t mc1, int64_t nc0, int64_t nc1) { + int tid = get_thread_num(); + float* y_buf = y_buffer.data_ptr() + tid * block_size; + at::BFloat16 *dqA_buffer = nullptr, *dqB_buffer = nullptr; + float* ukernel_buf = nullptr; +#if defined(CPU_CAPABILITY_AVX512) + at::BFloat16* micro_gemm_buf = micro_gemm_buffer.data_ptr() + tid * buffer_size; + ukernel_buf = reinterpret_cast(micro_gemm_buf); +#ifndef CPUBLAS_BRGEMM_F8F8F32 + dqA_buffer = micro_gemm_buf; + dqB_buffer = micro_gemm_buf + block_m * block_k; + ukernel_buf = reinterpret_cast(micro_gemm_buf + block_m * block_k + block_k * block_n); +#endif +#endif + + loop_2d(mc0, mc1, nc0, nc1, block_n * K, [&](int64_t mci, int64_t nc, int64_t) { + int64_t m_size = mci * block_m + block_m > M ? M - mci * block_m : block_m; + zero_buffer(y_buf, m_size * block_n); + for (int kci = 0; kci < Kc; ++kci) { + auto scales_a = a_scales_ptr + mci * block_m * num_groups + kci / block_per_group; + auto scales_b = b_scales_ptr + nc * block_n * num_groups + kci / block_per_group * block_n; + _micro_gemm( + /* C */ y_buf, + /* A */ a_ptr + mci * block_m * K + kci * block_k, + /* scales_a */ scales_a, + /* B */ b_ptr + (nc * Kc + kci) * block_n * block_k, + /* scales_b */ scales_b, + /* M */ m_size, + /* K */ block_k, + /* lda */ K, + /* ldc */ block_n, + /* ldsa */ ldsa, + /* ukernel_buf */ ukernel_buf, + /* dqA_buf */ dqA_buffer, + /* dqB_buf */ dqB_buffer); + } + auto scales_a = act_quant_mode == PER_TENSOR ? a_scales_ptr + : act_quant_mode == PER_ROW ? a_scales_ptr + mci * block_m + : nullptr; + auto scales_b = wei_quant_mode == PER_TENSOR ? b_scales_ptr + : wei_quant_mode == PER_ROW ? b_scales_ptr + nc * block_n + : nullptr; + auto bias_data = bias_ptr ? bias_ptr + nc * block_n : nullptr; + store_out( + y_buf, c_ptr + mci * block_m * N + nc * block_n, m_size, N /*lda*/, scales_a, scales_b, bias_data); + }); + + if constexpr (cpublas_can_pack) { + at::native::cpublas::brgemm_release(); + } + }); +} + +} // anonymous namespace + +void tinygemm_kernel( + float* C, + const at::Float8_e4m3fn* A, + const float* scales_a, + const at::Float8_e4m3fn* B, + const float* scales_b, + int64_t M, + int64_t K, + int64_t lda, + int64_t ldc, + int64_t ldsa, + float* ukernel_buf, + at::BFloat16* dqA_buf, + at::BFloat16* dqB_buf) { + // act_quant_mode = per row (2), wei_quant_mode = per group (3) + if (cpublas_could_pack()) { + _micro_gemm(C, A, scales_a, B, scales_b, M, K, lda, ldc, ldsa, ukernel_buf, dqA_buf, dqB_buf); + } else { + _micro_gemm(C, A, scales_a, B, scales_b, M, K, lda, ldc, ldsa, ukernel_buf, dqA_buf, dqB_buf); + } +} + +#define INSTANTIATE_TINYGEMM_TEMPLATE() \ + void tinygemm_kernel( \ + float* C, \ + const at::Float8_e4m3fn* A, \ + const float* scales_a, \ + const at::Float8_e4m3fn* B, \ + const float* scales_b, \ + int64_t M, \ + int64_t K, \ + int64_t lda, \ + int64_t ldc, \ + int64_t ldsa, \ + float* ukernel_buf, \ + at::BFloat16* dqA_buf, \ + at::BFloat16* dqB_buf) + +INSTANTIATE_TINYGEMM_TEMPLATE(); +/* +return: packed_weight, packed_scales +*/ +std::tuple float8_linear_prepack_impl(const at::Tensor& weight, const at::Tensor& scales) { + // weight shape = [N, K] + // scales shape = [N, G] + TORCH_CHECK(weight.dim() == 2, "Float8 linear CPU: Weight should be a 2D tensor for packing"); + TORCH_CHECK(weight.size(1) % 2 == 0, "Float8 linear CPU: Weight should have even number of columns for packing"); + + auto new_scales = scales; + if (new_scales.dim() == 1) { + new_scales.unsqueeze_(1); + } + new_scales = new_scales.to(at::kFloat); + int N = weight.size(0); + int K = weight.size(1); + int G = scales.size(1); + int group_size = K / G; + int block_k = group_size > 128 ? 128 : group_size; + while (K % block_k != 0) { + block_k /= 2; + } + TORCH_CHECK( + block_k > 0 && block_k <= group_size, "Float8 linear CPU: Invalid block_k size, should be in (0, group_size]"); + constexpr int block_n = BLOCK_N; + int Nc = N / block_n; + int Kc = K / block_k; + + // Reorder weight to [N/block_n, K/block_k, block_k, block_n] + // Reorder scales to [N/block_n, G, block_n] + auto weight_view = weight.view({Nc, block_n, Kc, block_k}); + at::Tensor weight_reordered = weight_view.permute({0, 2, 3, 1}).contiguous(); + at::Tensor blocked_weight; + at::Tensor blocked_scales = new_scales.view({Nc, block_n, G}).permute({0, 2, 1}).contiguous(); + +#if defined(CPU_CAPABILITY_AVX512) + if (cpublas_could_pack()) { +#ifdef CPUBLAS_BRGEMM_F8F8F32 + constexpr int vnni_size = 4; // for fp8 +#else + constexpr int vnni_size = 2; // for float16 +#endif + blocked_weight = at::empty({Nc, Kc, block_k, block_n}, weight.options()); + auto weight_ptr = reinterpret_cast(weight_reordered.data_ptr()); + auto blocked_weight_ptr = reinterpret_cast(blocked_weight.data_ptr()); + int64_t num_blocks = Nc * Kc; + at::parallel_for(0, num_blocks, 1, [&](int64_t begin, int64_t end) { + for (const auto i : c10::irange(begin, end)) { + auto in_ptr = weight_ptr + i * block_k * block_n; + auto out_ptr = blocked_weight_ptr + i * block_k * block_n; + + // Reorder weight block to VNNI + // plain shape = [block_k, block_n] + // packed shape = [block_k / VNNI_SIZE, block_n, VNNI_SIZE] viewed as [block_k, block_n] + constexpr int n_group_size = 8; + constexpr int n_group = block_n / n_group_size; // 4 + for (int nb = 0; nb < n_group; ++nb) { + for (int k = 0; k < block_k; k += vnni_size) { + for (int ni = 0; ni < n_group_size; ++ni) { + for (int ki = 0; ki < vnni_size; ++ki) { + int src_idx = nb * n_group_size + ni + (k + ki) * block_n; + int dst_idx = (nb * n_group_size + ni) * vnni_size + k * block_n + ki; + *(out_ptr + dst_idx) = *(in_ptr + src_idx); + } + } + } + } + } + }); + } else +#endif + { + blocked_weight = weight_reordered; + } + + return std::make_tuple(std::move(blocked_weight), std::move(blocked_scales)); +} + +// AVX512 optimized channel-wise max computation +inline void +compute_channel_max_avx512(const float* data, float* max_vals, int64_t num_channels, int64_t elements_per_channel) { + at::parallel_for(0, num_channels, 1, [&](int64_t start, int64_t end) { + for (int64_t c = start; c < end; ++c) { + const float* channel_data = data + c * elements_per_channel; + float max_val = 0.0f; + + int64_t i = 0; + // Process 16 elements at a time + __m512 max_vec = _mm512_setzero_ps(); + + for (; i <= elements_per_channel - 16; i += 16) { + __m512 data_vec = _mm512_loadu_ps(&channel_data[i]); + __m512 abs_vec = _mm512_abs_ps(data_vec); + max_vec = _mm512_max_ps(max_vec, abs_vec); + } + + // Horizontal max of the vector + float hmax = _mm512_reduce_max_ps(max_vec); + max_val = std::max(max_val, hmax); + + // Handle remaining elements + for (; i < elements_per_channel; ++i) { + max_val = std::max(max_val, std::abs(channel_data[i])); + } + + max_vals[c] = max_val; + } + }); +} + +// AVX512 optimized scaling and clamping for channel-wise quantization +inline void scale_clamp_channelwise_avx512( + const float* src, + float* dst, + const float* scales, + int64_t num_channels, + int64_t elements_per_channel, + float quant_max, + float neg_quant_max) { + at::parallel_for(0, num_channels, 1, [&](int64_t start, int64_t end) { + for (int64_t c = start; c < end; ++c) { + float scale_val = scales[c]; + float scale_reciprocal = 1.0f / scale_val; + + const float* channel_src = src + c * elements_per_channel; + float* channel_dst = dst + c * elements_per_channel; + + int64_t i = 0; + const __m512 scale_recip_vec = _mm512_set1_ps(scale_reciprocal); + const __m512 quant_max_vec = _mm512_set1_ps(quant_max); + const __m512 neg_quant_max_vec = _mm512_set1_ps(neg_quant_max); + + for (; i <= elements_per_channel - 16; i += 16) { + __m512 src_vec = _mm512_loadu_ps(&channel_src[i]); + __m512 scaled_vec = _mm512_mul_ps(src_vec, scale_recip_vec); + __m512 clamped_vec = _mm512_min_ps(_mm512_max_ps(scaled_vec, neg_quant_max_vec), quant_max_vec); + _mm512_storeu_ps(&channel_dst[i], clamped_vec); + } + + // Handle remaining elements + for (; i < elements_per_channel; ++i) { + float scaled = channel_src[i] * scale_reciprocal; + channel_dst[i] = std::clamp(scaled, neg_quant_max, quant_max); + } + } + }); +} + +// AVX512 optimized scaling and clamping for global quantization +inline void scale_clamp_global_avx512( + const float* src, float* dst, int64_t size, float scale_reciprocal, float quant_max, float neg_quant_max) { + if (size <= 0) return; + + const __m512 scale_recip_vec = _mm512_set1_ps(scale_reciprocal); + const __m512 quant_max_vec = _mm512_set1_ps(quant_max); + const __m512 neg_quant_max_vec = _mm512_set1_ps(neg_quant_max); + + at::parallel_for(0, size, 4096, [&](int64_t start, int64_t end) { + int64_t block_size = end - start; + const float* block_src = src + start; + float* block_dst = dst + start; + + int64_t i = 0; + for (; i <= block_size - 16; i += 16) { + __m512 src_vec = _mm512_loadu_ps(&block_src[i]); + __m512 scaled_vec = _mm512_mul_ps(src_vec, scale_recip_vec); + __m512 clamped_vec = _mm512_min_ps(_mm512_max_ps(scaled_vec, neg_quant_max_vec), quant_max_vec); + _mm512_storeu_ps(&block_dst[i], clamped_vec); + } + + // Handle remaining elements in block + for (; i < block_size; ++i) { + float scaled = block_src[i] * scale_reciprocal; + block_dst[i] = std::clamp(scaled, neg_quant_max, quant_max); + } + }); +} + +std::tuple +_quantize_fp8e4m3(const at::Tensor& t, bool channelwise, c10::optional scale_opt = c10::nullopt) { + constexpr float quant_max = 448.0f; // torch.finfo(torch.float8_e4m3fn).max + constexpr float eps = std::numeric_limits::epsilon(); + + // Ensure input is contiguous and in float32 + auto t_float = t.to(at::ScalarType::Float).contiguous(); + at::Tensor qt; + at::Tensor scale_tensor; + + if (channelwise) { + // Channel-wise quantization with AVX512 optimization + int64_t num_channels = t_float.size(0); + int64_t elements_per_channel = t_float.numel() / num_channels; + + if (elements_per_channel * num_channels != t_float.numel()) { + throw std::runtime_error("Tensor must be divisible by number of channels for channel-wise quantization"); + } + + // Allocate scale tensor + scale_tensor = at::empty({num_channels}, t_float.options()); + float* scale_data = scale_tensor.data_ptr(); + + // Compute channel-wise max using AVX512 + compute_channel_max_avx512(t_float.data_ptr(), scale_data, num_channels, elements_per_channel); + + // Apply quant_max and EPS + at::parallel_for(0, num_channels, 1, [&](int64_t start, int64_t end) { + for (int64_t c = start; c < end; ++c) { + scale_data[c] = std::max(scale_data[c] / quant_max, eps); + } + }); + + // Create output tensor for quantized values + qt = at::empty_like(t_float); + + // Scale and clamp using AVX512 + scale_clamp_channelwise_avx512( + t_float.data_ptr(), + qt.data_ptr(), + scale_data, + num_channels, + elements_per_channel, + quant_max, + -quant_max); + } else { + // Global quantization with AVX512 optimization + if (!scale_opt.has_value()) { + throw std::runtime_error("Scale must be provided for non-channelwise quantization in AVX512 version"); + } + + scale_tensor = scale_opt.value().to(at::ScalarType::Float).contiguous(); + + // Handle scalar scale case + float scale_val; + if (scale_tensor.numel() == 1) { + scale_val = std::max(scale_tensor.item(), eps); + } else { + // Compute max of scale tensor if it's not scalar + float* scale_data = scale_tensor.data_ptr(); + __m512 max_vec = _mm512_set1_ps(eps); + + int64_t i = 0; + for (; i <= scale_tensor.numel() - 16; i += 16) { + __m512 scale_vec = _mm512_loadu_ps(&scale_data[i]); + max_vec = _mm512_max_ps(max_vec, scale_vec); + } + + float hmax = _mm512_reduce_max_ps(max_vec); + for (; i < scale_tensor.numel(); ++i) { + hmax = std::max(hmax, scale_data[i]); + } + scale_val = std::max(hmax, eps); + } + + scale_tensor = at::tensor({scale_val}, scale_tensor.options()); + float scale_reciprocal = 1.0f / scale_val; + + // Create output tensor + qt = at::empty_like(t_float); + + // Scale and clamp using AVX512 + scale_clamp_global_avx512( + t_float.data_ptr(), qt.data_ptr(), t_float.numel(), scale_reciprocal, quant_max, -quant_max); + } + + // Final conversion to FP8 E4M3FN using ATen's native conversion + qt = at::_to_copy(qt, at::ScalarType::Float8_e4m3fn); + + return std::make_tuple(qt, scale_tensor); +} + +inline __m128i cvtfp32_fp8e4m3(__m512& src) { + // cvt 16x32 from fp32 to fp8 e4m3 + const __m512i sign_mask = _mm512_set1_epi32(0x80000000); + const __m512i fp8_max = _mm512_set1_epi32(UINT32_C(1087) << 20); + const __m512i denorm_thresh = _mm512_set1_epi32(UINT32_C(121) << 23); + const __m512i denorm_mask = _mm512_set1_epi32(UINT32_C(141) << 23); + const __m512i bias_part1 = _mm512_set1_epi32((uint32_t)(7 - 127) << 23); + const __m512i rounding_bias = _mm512_set1_epi32(0x7FFFF); + __m512i f_bits = _mm512_castps_si512(src); + // Extract and save sign + __m512i sign = _mm512_and_epi32(f_bits, sign_mask); + f_bits = _mm512_xor_epi32(f_bits, sign); + + // Prepare result containers + __m512i result = _mm512_setzero_si512(); + + // Step 1: Handle case of overflow + // (f_bits >= fp8_max): set result = 0x7f + __mmask16 overflow_mask = _mm512_cmpge_epu32_mask(f_bits, fp8_max); + if (overflow_mask) { + result = _mm512_mask_set1_epi32(result, overflow_mask, 0x7f); + } + + // Step 2: Handle small numbers (denormals) + // Small numbers (f_bits < denorm_thresh) + __mmask16 denorm_thresh_mask = _mm512_cmplt_epu32_mask(f_bits, denorm_thresh); + + if (denorm_thresh_mask) { + __m512 small_input = _mm512_castsi512_ps(f_bits); + __m512 small_denorm = _mm512_add_ps(small_input, _mm512_castsi512_ps(denorm_mask)); + __m512i small_denorm_bits = _mm512_castps_si512(small_denorm); + __m512i small_result = _mm512_sub_epi32(small_denorm_bits, denorm_mask); + result = _mm512_mask_mov_epi32(result, denorm_thresh_mask, small_result); + } + + // Step 3: Handle normal numbers + __mmask16 normal_mask = ~(overflow_mask | denorm_thresh_mask); + + if (normal_mask) { + // mant_odd = (f_bits >> 20) & 1 + __m512i mant_odd = _mm512_and_epi32(_mm512_srli_epi32(f_bits, 20), _mm512_set1_epi32(1)); + // f_bits += bias_part1 + rounding_bias + __m512i rounded = _mm512_add_epi32(f_bits, bias_part1); + rounded = _mm512_add_epi32(rounded, rounding_bias); + // Add mant_odd + rounded = _mm512_add_epi32(rounded, mant_odd); + // Shift right by 20 bits + __m512i normal_result = _mm512_srli_epi32(rounded, 20); + result = _mm512_mask_mov_epi32(result, normal_mask, normal_result); + } + + // Merge back the sign + __m512i sign_shifted = _mm512_srli_epi32(sign, 24); + result = _mm512_or_epi32(result, sign_shifted); + + // Now result is 16 x 32-bit integers, but we only need 8-bit for each + __m512i packed = _mm512_and_si512(result, _mm512_set1_epi32(0xFF)); + + // Narrow 32-bit integers to 8-bit + return _mm512_cvtepi32_epi8(packed); +} + +__attribute__((target("avx10.2"))) inline __m128i cvtfp32_fp8e4m3_avx10_2(__m512& src) { + __m256i f16_vec = _mm512_cvt_roundps_ph(src, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC); + return _mm256_cvtph_hf8(_mm256_castsi256_ph(f16_vec)); +} + +// Carries the whole loop rather than just the conversion, as GCC cannot inline an avx10.2 +// function into a caller lacking that target, so branching per vector is costly. +__attribute__((target("avx10.2"))) void quantize_bf16_block_avx10_2( + const at::BFloat16* __restrict__ src, at::Float8_e4m3fn* __restrict__ dst, int64_t count, float scale_reciprocal) { + constexpr float quant_max = 448.0f; // torch.finfo(torch.float8_e4m3fn).max + const __m512 scale_recip_vec = _mm512_set1_ps(scale_reciprocal); + const __m512 quant_max_vec = _mm512_set1_ps(quant_max); + const __m512 neg_quant_max_vec = _mm512_set1_ps(-quant_max); + + for (int64_t i = 0; i <= count - 32; i += 32) { + __m256i src_vec1 = _mm256_loadu_si256(reinterpret_cast(&src[i])); + __m256i src_vec2 = _mm256_loadu_si256(reinterpret_cast(&src[i + 16])); + + __m512 fp32_vec1 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(src_vec1), 16)); + __m512 fp32_vec2 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(src_vec2), 16)); + + __m512 scaled_vec1 = _mm512_mul_ps(fp32_vec1, scale_recip_vec); + __m512 clamped_vec1 = _mm512_min_ps(_mm512_max_ps(scaled_vec1, neg_quant_max_vec), quant_max_vec); + + __m512 scaled_vec2 = _mm512_mul_ps(fp32_vec2, scale_recip_vec); + __m512 clamped_vec2 = _mm512_min_ps(_mm512_max_ps(scaled_vec2, neg_quant_max_vec), quant_max_vec); + + __m128i fp8_vec1 = cvtfp32_fp8e4m3_avx10_2(clamped_vec1); + __m128i fp8_vec2 = cvtfp32_fp8e4m3_avx10_2(clamped_vec2); + _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + i), fp8_vec1); + _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + i + 16), fp8_vec2); + } +} + +std::tuple _quantize_fp8e4m3_bf16_per_tensor_no_scale(const at::Tensor& t) { + constexpr float quant_max = 448.0f; // torch.finfo(torch.float8_e4m3fn).max + constexpr float eps = std::numeric_limits::epsilon(); + + // Input validation and preparation + assert(t.scalar_type() == at::ScalarType::BFloat16); + auto t_bf16 = t.contiguous(); + int64_t num_channels = t_bf16.size(0); + int64_t elements_per_channel = t_bf16.numel() / num_channels; + assert(elements_per_channel % 32 == 0); // do not consider tile currently + at::Tensor quant_t = at::empty(t_bf16.sizes(), t_bf16.options().dtype(at::kFloat8_e4m3fn)); + + // Allocate output tensors + at::Tensor scale_tensor = at::empty({num_channels}, t_bf16.options().dtype(at::ScalarType::Float)); + float* scale_data = scale_tensor.data_ptr(); + + // Unified processing: compute max, apply scale, and quantize in single pass + at::parallel_for(0, num_channels, 1, [&](int64_t start, int64_t end) { + for (int64_t c = start; c < end; ++c) { + const at::BFloat16* channel_src = t_bf16.data_ptr() + c * elements_per_channel; + at::Float8_e4m3fn* quant_dst = quant_t.data_ptr() + c * elements_per_channel; + + // Step 1: Compute channel-wise max using AVX512 + float channel_max = 0.0f; + int64_t i = 0; + + // Process 32 elements at a time for max computation + for (; i <= elements_per_channel - 32; i += 32) { + // Load 32 BF16 values (2x 256-bit vectors) + __m256i src_vec1 = _mm256_loadu_si256(reinterpret_cast(&channel_src[i])); + __m256i src_vec2 = _mm256_loadu_si256(reinterpret_cast(&channel_src[i + 16])); + + // Convert BF16 to FP32 + __m512i fp32_int1 = _mm512_cvtepu16_epi32(src_vec1); + __m512 fp32_vec1 = _mm512_castsi512_ps(_mm512_slli_epi32(fp32_int1, 16)); + + __m512i fp32_int2 = _mm512_cvtepu16_epi32(src_vec2); + __m512 fp32_vec2 = _mm512_castsi512_ps(_mm512_slli_epi32(fp32_int2, 16)); + + // Compute absolute values + __m512 abs_vec1 = _mm512_abs_ps(fp32_vec1); + __m512 abs_vec2 = _mm512_abs_ps(fp32_vec2); + + // Find max in each vector + float max1 = _mm512_reduce_max_ps(abs_vec1); + float max2 = _mm512_reduce_max_ps(abs_vec2); + channel_max = std::max(channel_max, std::max(max1, max2)); + } + + // Step 2: Apply quant_max and EPS to compute scale + float scale_val = std::max(channel_max / quant_max, eps); + float scale_reciprocal = 1.0f / scale_val; + scale_data[c] = scale_val; + + // Step 3: Scale and clamp using AVX512 (reuse the same loop structure) + if (avx10_2_available()) { + quantize_bf16_block_avx10_2(channel_src, quant_dst, elements_per_channel, scale_reciprocal); + continue; + } + i = 0; + const __m512 scale_recip_vec = _mm512_set1_ps(scale_reciprocal); + const __m512 quant_max_vec = _mm512_set1_ps(quant_max); + const __m512 neg_quant_max_vec = _mm512_set1_ps(-quant_max); + + for (; i <= elements_per_channel - 32; i += 32) { + // Load 32 BF16 values + __m256i src_vec1 = _mm256_loadu_si256(reinterpret_cast(&channel_src[i])); + __m256i src_vec2 = _mm256_loadu_si256(reinterpret_cast(&channel_src[i + 16])); + + // Convert BF16 to FP32 + __m512 fp32_vec1 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(src_vec1), 16)); + __m512 fp32_vec2 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(src_vec2), 16)); + + // Scale and clamp + __m512 scaled_vec1 = _mm512_mul_ps(fp32_vec1, scale_recip_vec); + __m512 clamped_vec1 = _mm512_min_ps(_mm512_max_ps(scaled_vec1, neg_quant_max_vec), quant_max_vec); + + __m512 scaled_vec2 = _mm512_mul_ps(fp32_vec2, scale_recip_vec); + __m512 clamped_vec2 = _mm512_min_ps(_mm512_max_ps(scaled_vec2, neg_quant_max_vec), quant_max_vec); + + __m128i fp8_vec1 = cvtfp32_fp8e4m3(clamped_vec1); + __m128i fp8_vec2 = cvtfp32_fp8e4m3(clamped_vec2); + _mm_storeu_si128(reinterpret_cast<__m128i*>(quant_dst + i), fp8_vec1); + _mm_storeu_si128(reinterpret_cast<__m128i*>(quant_dst + i + 16), fp8_vec2); + } + } + }); + + return std::make_tuple(quant_t, scale_tensor); +} + +std::tuple +_quantize_fp8e4m3_bf16_per_tensor_with_scale(const at::Tensor& t, at::Tensor& scale_tensor) { + constexpr float quant_max = 448.0f; // torch.finfo(torch.float8_e4m3fn).max + constexpr float eps = std::numeric_limits::epsilon(); + + // Input validation + assert(t.scalar_type() == at::ScalarType::BFloat16); + + auto t_bf16 = t.contiguous(); + auto scale_tensor_contig = scale_tensor.contiguous().to(at::ScalarType::Float); + + // Get scale value (handle scalar or tensor) + float scale_val; + if (scale_tensor_contig.numel() == 1) { + scale_val = scale_tensor_contig.item(); + } else { + // Take max if scale is a tensor (though should be scalar for global quantization) + scale_val = scale_tensor_contig.max().item(); + } + + // Apply EPS to scale + scale_val = std::max(scale_val, eps); + float scale_reciprocal = 1.0f / scale_val; + + at::Tensor scale_output = at::tensor({scale_val}, scale_tensor_contig.options()); + + // Apply scale and clamp using AVX512 + const at::BFloat16* src_data = t_bf16.data_ptr(); + at::Tensor quant_t = at::empty(t_bf16.sizes(), t_bf16.options().dtype(at::kFloat8_e4m3fn)); + int64_t total_elements = t_bf16.numel(); + const __m512 scale_recip_vec = _mm512_set1_ps(scale_reciprocal); + const __m512 quant_max_vec = _mm512_set1_ps(quant_max); + const __m512 neg_quant_max_vec = _mm512_set1_ps(-quant_max); + int64_t num_channels = t_bf16.size(0); + int64_t elements_per_channel = total_elements / num_channels; + assert(elements_per_channel % 32 == 0); // do not consider tile currently + // Process in parallel blocks + at::parallel_for(0, num_channels, 1, [&](int64_t start, int64_t end) { + for (int64_t c = start; c < end; ++c) { + const at::BFloat16* src_data = t_bf16.data_ptr() + c * elements_per_channel; + at::Float8_e4m3fn* quant_t_data = quant_t.data_ptr() + c * elements_per_channel; + + if (avx10_2_available()) { + quantize_bf16_block_avx10_2(src_data, quant_t_data, elements_per_channel, scale_reciprocal); + continue; + } + + int64_t i = 0; + // Process 32 elements at a time using AVX512 + for (; i <= elements_per_channel - 32; i += 32) { + // Load 32 BF16 values (2x 256-bit vectors) + __m256i src_vec1 = _mm256_loadu_si256(reinterpret_cast(&src_data[i])); + __m256i src_vec2 = _mm256_loadu_si256(reinterpret_cast(&src_data[i + 16])); + + // Convert BF16 to FP32 + __m512 fp32_vec1 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(src_vec1), 16)); + __m512 fp32_vec2 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(src_vec2), 16)); + + // Scale and clamp + __m512 scaled_vec1 = _mm512_mul_ps(fp32_vec1, scale_recip_vec); + __m512 clamped_vec1 = _mm512_min_ps(_mm512_max_ps(scaled_vec1, neg_quant_max_vec), quant_max_vec); + + __m512 scaled_vec2 = _mm512_mul_ps(fp32_vec2, scale_recip_vec); + __m512 clamped_vec2 = _mm512_min_ps(_mm512_max_ps(scaled_vec2, neg_quant_max_vec), quant_max_vec); + + __m128i fp8_vec1 = cvtfp32_fp8e4m3(clamped_vec1); + __m128i fp8_vec2 = cvtfp32_fp8e4m3(clamped_vec2); + _mm_storeu_si128(reinterpret_cast<__m128i*>(quant_t_data + i), fp8_vec1); + _mm_storeu_si128(reinterpret_cast<__m128i*>(quant_t_data + i + 16), fp8_vec2); + } + } + }); + + return std::make_tuple(quant_t, scale_output); +} + +std::tuple +_quantize_fp8e4m3_vec(const at::Tensor& t, bool channelwise, c10::optional scale_opt) { + if (channelwise) { + return _quantize_fp8e4m3_bf16_per_tensor_no_scale(t); + } else { + assert(scale_opt.has_value()); + return _quantize_fp8e4m3_bf16_per_tensor_with_scale(t, scale_opt.value()); + } +} + +// Public wrapper for torch op registration +std::tuple +quantize_fp8e4m3_vec(const at::Tensor& t, bool channelwise, c10::optional scale_opt) { + return _quantize_fp8e4m3_vec(t, channelwise, scale_opt); +} + +at::Tensor fp8_scaled_mm_with_quant( + const at::Tensor& act, + const std::optional& act_scales, + bool channelwise, + const at::Tensor& weight, + const at::Tensor& weight_scales, + const std::optional& bias, + at::ScalarType output_dtype) { + // TODO: refine here to use tensor ptr and in below dispatch api. + std::tuple quant_act = act.scalar_type() == at::ScalarType::BFloat16 + ? _quantize_fp8e4m3_vec(act, channelwise, act_scales) + : _quantize_fp8e4m3(act, channelwise, act_scales); + auto input = std::get<0>(quant_act); + auto input_scales = std::get<1>(quant_act); + int64_t N = weight.dim() == 4 ? weight.size(0) * weight.size(-1) : weight.size(0); + int act_quant_mode = input_scales.numel() == 1 ? PER_TENSOR + : input_scales.numel() == input.numel() / input.size(-1) ? PER_ROW + : PER_GROUP; + int wei_quant_mode = weight_scales.numel() == 1 ? PER_TENSOR : weight_scales.numel() == N ? PER_ROW : PER_GROUP; + // Case to fall back + if (weight.dim() == 2) { + TORCH_CHECK( + act_quant_mode != PER_GROUP && wei_quant_mode != PER_GROUP, + "FP8 linear: Per-group quantization is not supported in the fallback path"); + auto y_fp32 = at::linear(input.to(at::kFloat).mul_(input_scales), weight.to(at::kFloat).mul_(weight_scales), bias); + return y_fp32.to(output_dtype); + } + + static bool cpublas_can_pack = cpublas_could_pack(); + auto out_sizes = input.sizes().vec(); + out_sizes.back() = N; + auto output = at::empty(out_sizes, input.options().dtype(output_dtype)); + +#define AT_DISPATCH_FP8_LINEAR_KERNEL(OUT_DTYPE, CAN_PACK, A_QUANT_MODE, B_QUANT_MODE, ...) \ + AT_DISPATCH_BOOL_NO_RETURN( \ + CAN_PACK, \ + "cpublas_can_pack", \ + can_pack, \ + AT_DISPATCH_QUANT_MODE_NO_RETURN( \ + A_QUANT_MODE, \ + "act_quant_mode", \ + a_quant_mode, \ + AT_DISPATCH_QUANT_MODE_NO_RETURN( \ + B_QUANT_MODE, \ + "wei_quant_mode", \ + b_quant_mode, \ + AT_DISPATCH_OUT_TYPES(OUT_DTYPE, "out_dtype", __VA_ARGS__)))) + + AT_DISPATCH_FP8_LINEAR_KERNEL(output_dtype, cpublas_can_pack, act_quant_mode, wei_quant_mode, [&]() { + _float8_linear_impl( + input, input_scales, weight, weight_scales, bias, output); + }); + return output; +} + +at::Tensor float8_linear_impl( + const at::Tensor& input, + const at::Tensor& input_scales, + const at::Tensor& weight, + const at::Tensor& weight_scales, + const std::optional& bias, + at::ScalarType output_dtype) { + int64_t N = weight.dim() == 4 ? weight.size(0) * weight.size(-1) : weight.size(0); + int act_quant_mode = input_scales.numel() == 1 ? PER_TENSOR + : input_scales.numel() == input.numel() / input.size(-1) ? PER_ROW + : PER_GROUP; + int wei_quant_mode = weight_scales.numel() == 1 ? PER_TENSOR : weight_scales.numel() == N ? PER_ROW : PER_GROUP; + // Case to fall back + if (weight.dim() == 2) { + TORCH_CHECK( + act_quant_mode != PER_GROUP && wei_quant_mode != PER_GROUP, + "FP8 linear: Per-group quantization is not supported in the fallback path"); + auto y_fp32 = at::linear(input.to(at::kFloat).mul_(input_scales), weight.to(at::kFloat).mul_(weight_scales), bias); + return y_fp32.to(output_dtype); + } + + static bool cpublas_can_pack = cpublas_could_pack(); + auto out_sizes = input.sizes().vec(); + out_sizes.back() = N; + auto output = at::empty(out_sizes, input.options().dtype(output_dtype)); + +#define AT_DISPATCH_FP8_LINEAR_KERNEL(OUT_DTYPE, CAN_PACK, A_QUANT_MODE, B_QUANT_MODE, ...) \ + AT_DISPATCH_BOOL_NO_RETURN( \ + CAN_PACK, \ + "cpublas_can_pack", \ + can_pack, \ + AT_DISPATCH_QUANT_MODE_NO_RETURN( \ + A_QUANT_MODE, \ + "act_quant_mode", \ + a_quant_mode, \ + AT_DISPATCH_QUANT_MODE_NO_RETURN( \ + B_QUANT_MODE, \ + "wei_quant_mode", \ + b_quant_mode, \ + AT_DISPATCH_OUT_TYPES(OUT_DTYPE, "out_dtype", __VA_ARGS__)))) + + AT_DISPATCH_FP8_LINEAR_KERNEL(output_dtype, cpublas_can_pack, act_quant_mode, wei_quant_mode, [&]() { + _float8_linear_impl( + input, input_scales, weight, weight_scales, bias, output); + }); + return output; +} \ No newline at end of file diff --git a/csrc/cpu/sgl-kernels/moe.cpp b/csrc/cpu/sgl-kernels/moe.cpp index 3afdc1bd67..09a170eb6d 100644 --- a/csrc/cpu/sgl-kernels/moe.cpp +++ b/csrc/cpu/sgl-kernels/moe.cpp @@ -836,6 +836,11 @@ static inline void check_moe_scales( TORCH_CHECK(w2_scale.has_value(), "missing w2_scale for mxfp4."); TORCH_CHECK(w1_scale.value().scalar_type() == at::kByte, "expect w1_scale to be uint8."); TORCH_CHECK(w2_scale.value().scalar_type() == at::kByte, "expect w2_scale to be uint8."); + } else if constexpr (quant == CPUQuantMethod::FP8_W8A8) { + TORCH_CHECK(w1_scale.has_value(), "missing w1_scale for fp8 w8a8."); + TORCH_CHECK(w2_scale.has_value(), "missing w2_scale for fp8 w8a8."); + TORCH_CHECK(block_size.has_value(), "missing block_size for fp8 w8a8."); + TORCH_CHECK(block_size.value().size() == 2, "expect block_size.size() to be 2."); } } @@ -850,9 +855,14 @@ static inline void check_moe_scales( check_moe_scales(w1_scale, w2_scale, block_size); } else if (moe_comp_method == CPUQuantMethod::MXFP4) { check_moe_scales(w1_scale, w2_scale, block_size); + } else if (moe_comp_method == CPUQuantMethod::FP8_W8A8) { + check_moe_scales(w1_scale, w2_scale, block_size); } } +// CHECK_MOE_SCALES_FP8: validate scale shapes for FP8_W8A16. +// w1s shape [E, 2N/block_N, K/block_K, block_N]: DIM0=1, DIM1=2 +// w2s shape [E, K/block_N, N/block_K]: DIM0=1, DIM1=2 #define CHECK_MOE_SCALES_FP8(DIM0, DIM1) \ auto w1s = w1_scale.value(); \ auto w2s = w2_scale.value(); \ @@ -864,6 +874,17 @@ static inline void check_moe_scales( TORCH_CHECK(w2s.size(DIM0) == div_up(K, block_size_N)); \ TORCH_CHECK(w2s.size(DIM1) == div_up(N, block_size_K)) +// CHECK_MOE_SCALES_FP8_W8A8: extract scale info for FP8_W8A8 without shape +// validation, because w1s is pre-packed by float8_linear_prepack_cpu with +// shape [E, 2N/BLOCK_N, G, BLOCK_N] where BLOCK_N=32 != block_size_N(128). +// This matches the sglang upstream behavior (checks are commented out there). +#define CHECK_MOE_SCALES_FP8_W8A8(DIM0, DIM1) \ + auto w1s = w1_scale.value(); \ + auto w2s = w2_scale.value(); \ + auto block_size_val = block_size.value(); \ + int64_t block_size_N = block_size_val[0]; \ + int64_t block_size_K = block_size_val[1] + // hidden_states: [M, K] // w1: [E, 2N, K] or [E, 2N, K / 2] for uint8 // w2: [E, K, N] or [E, K, N / 2] for uint8 @@ -883,6 +904,7 @@ void fused_experts_cpu( const std::optional& w2_scale, const std::optional& w1_zero, const std::optional& w2_zero, + const std::optional& a1_scale, const std::optional> block_size, const std::optional& w1_bias, const std::optional& w2_bias, @@ -895,7 +917,10 @@ void fused_experts_cpu( constexpr int64_t BLOCK_M = block_size_m(); constexpr int64_t BLOCK_N = block_size_n(); - const auto st = hidden_states.scalar_type(); + // FP8_W8A8 hidden_states input is FP8; scalar_t should be BF16 for output + const auto st = (moe_comp_method == CPUQuantMethod::FP8_W8A8) + ? at::kBFloat16 + : hidden_states.scalar_type(); // TODO: fused_topk_torch_native (CPU fallback for models like MiniMax) // returns int64 topk_ids; fused_experts_cpu requires int32. Remove the typecast after topk kernel is provided auto topk_ids_ = topk_ids.scalar_type() == at::kInt ? topk_ids : topk_ids.to(at::kInt); @@ -908,6 +933,9 @@ void fused_experts_cpu( if (moe_comp_method == CPUQuantMethod::INT4_W4A8 && is_vnni) { CHECK_DIM(4, w1); CHECK_DIM(4, w2); + } else if (moe_comp_method == CPUQuantMethod::FP8_W8A8) { + // w1 is prepackaged by float8_linear_prepack_cpu (5D), skip w1 dim check + CHECK_DIM(3, w2); } else { CHECK_DIM(3, w1); CHECK_DIM(3, w2); @@ -924,8 +952,9 @@ void fused_experts_cpu( int64_t M = hidden_states.size(0); int64_t K = hidden_states.size(1); - int64_t N = moe_comp_method == CPUQuantMethod::INT4_W4A8 ? w1_scale.value().size(1) * w1_scale.value().size(3) / 2 - : w1.size(1) / 2; + int64_t N = moe_comp_method == CPUQuantMethod::INT4_W4A8 + ? w1_scale.value().size(1) * w1_scale.value().size(3) / 2 + : (moe_comp_method == CPUQuantMethod::FP8_W8A8 ? w2.size(2) : w1.size(1) / 2); int64_t E = w1.size(0); int64_t topk = topk_weights_.size(1); @@ -935,7 +964,10 @@ void fused_experts_cpu( // check weight shapes CHECK_EQ(w2.size(0), E); - if (!(moe_comp_method == CPUQuantMethod::INT4_W4A8)) { + if (moe_comp_method == CPUQuantMethod::FP8_W8A8) { + CHECK_EQ(w2.size(1), K); + CHECK_EQ(packed_w2.size(2), packed_N); + } else if (!(moe_comp_method == CPUQuantMethod::INT4_W4A8)) { CHECK_EQ(w2.size(1), K); CHECK_EQ(packed_w1.size(2), packed_K / (moe_comp_method == CPUQuantMethod::INT4_W4A8 ? 2 : 1)); CHECK_EQ(packed_w2.size(2), packed_N / (moe_comp_method == CPUQuantMethod::INT4_W4A8 ? 2 : 1)); @@ -993,22 +1025,35 @@ void fused_experts_cpu( // 5. Aq_tmp : [M, K] or [M * topk, N] // 6. As_tmp : [M * topk] // - // for fp8 w8a16 and mxfp4: + // for fp8 w8a16, mxfp4, and fp8 w8a8: // 7. intermediate_cache0 : [M * topk, 2N] // 8. B_tmp : [T, MAX_CACHE_BLOCK_SIZE, BLOCK_N, std::max(K, N)] // + // for fp8 w8a8 additionally: + // 9. Ukernel_tmp : [T, 2 * BLOCK_M * BLOCK_N] (float32 for FP8×FP8 brgemm) + // int64_t buffer_size_nbytes = M * topk * N * 2 + M * topk * K * 2 + num_threads * BLOCK_M * K * - (moe_comp_method == CPUQuantMethod::INT8_W8A8 | moe_comp_method == CPUQuantMethod::INT4_W4A8 ? 1 : 2) + + (moe_comp_method == CPUQuantMethod::INT8_W8A8 | moe_comp_method == CPUQuantMethod::INT4_W4A8 + | moe_comp_method == CPUQuantMethod::FP8_W8A8 ? 1 : 2) + num_threads * 2 * BLOCK_M * BLOCK_N * sizeof(float); if (moe_comp_method == CPUQuantMethod::INT8_W8A8) { buffer_size_nbytes += std::max(M * K, M * topk * N) + M * topk * sizeof(float); } - if (moe_comp_method == CPUQuantMethod::FP8_W8A16 || moe_comp_method == CPUQuantMethod::MXFP4) { + if (moe_comp_method == CPUQuantMethod::FP8_W8A16 || moe_comp_method == CPUQuantMethod::MXFP4 + || moe_comp_method == CPUQuantMethod::FP8_W8A8) { buffer_size_nbytes += M * topk * 2 * N * 2 + num_threads * MAX_CACHE_BLOCK_SIZE * BLOCK_N * std::max(K, N) * 2; } + if (moe_comp_method == CPUQuantMethod::FP8_W8A8) { + // Ukernel buffer: float32 workspace per thread for FP8×FP8 brgemm + buffer_size_nbytes += num_threads * 2 * BLOCK_M * BLOCK_N * sizeof(float); +#ifndef CPUBLAS_BRGEMM_F8F8F32 + // BF16 fallback: also need dqA [T, BLOCK_M, BLOCK_K] + dqB [T, BLOCK_K, BLOCK_N] + buffer_size_nbytes += num_threads * (BLOCK_M * BLOCK_K + BLOCK_K * BLOCK_N) * 2; +#endif + } if (moe_comp_method == CPUQuantMethod::INT4_W4A8) { buffer_size_nbytes += M * topk * 2 * N * 2 + std::max(M * K, M * topk * N) + M * topk * sizeof(float) + num_threads * 2 * get_4bit_block_k_size(K / w1_scale.value().size(2)) * BLOCK_N; @@ -1181,6 +1226,47 @@ void fused_experts_cpu( E, topk, num_tokens_post_pad); + } else if (moe_comp_method == CPUQuantMethod::FP8_W8A8) { + // Stage 1: FP8×FP8 GEMM (hidden_states is FP8 input) + // Stage 2: BF16×FP8 GEMM (W8A16 path for down-projection) + at::Float8_e4m3fn* __restrict__ A_tmp = + (at::Float8_e4m3fn*)((void*)(intermediate_cache2 + M * topk * K)); + float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K)); + scalar_t* __restrict__ intermediate_cache0 = + (scalar_t*)((void*)(C_tmp + num_threads * 2 * BLOCK_M * BLOCK_N)); + scalar_t* __restrict__ B_tmp = + (scalar_t*)((void*)(intermediate_cache0 + M * topk * 2 * N)); + float* __restrict__ Ukernel_tmp = + (float*)((void*)(B_tmp + num_threads * MAX_CACHE_BLOCK_SIZE * BLOCK_N * std::max(K, N))); + + CHECK_MOE_SCALES_FP8_W8A8(1, 2); + fused_experts_fp8_a8_kernel_impl( + out_hidden_states.data_ptr(), + intermediate_cache0, + intermediate_cache1, + intermediate_cache2, + A_tmp, + B_tmp, + C_tmp, + Ukernel_tmp, + hidden_states.data_ptr(), + packed_w1.data_ptr(), + packed_w2.data_ptr(), + a1_scale.has_value() ? a1_scale.value().data_ptr() : nullptr, + w1s.data_ptr(), + w2s.data_ptr(), + block_size_N, + block_size_K, + topk_weights_.data_ptr(), + sorted_ids, + expert_ids, + offsets, + M, + N, + K, + E, + topk, + num_tokens_post_pad); } else { scalar_t* __restrict__ A_tmp = intermediate_cache2 + M * topk * K; float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K)); diff --git a/csrc/cpu/sgl-kernels/moe_fp8_w8a8.cpp b/csrc/cpu/sgl-kernels/moe_fp8_w8a8.cpp new file mode 100644 index 0000000000..0f393a34a8 --- /dev/null +++ b/csrc/cpu/sgl-kernels/moe_fp8_w8a8.cpp @@ -0,0 +1,432 @@ +// Adapted from +// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu + +// clang-format off + +#include + +#include "common.h" +#include "gemm.h" +#include "moe.h" + +namespace { + +// Byte-wise copy for FP8 elements (no remainder assumption) +inline void copy_stub2( + at::Float8_e4m3fn* __restrict__ out, + const at::Float8_e4m3fn* __restrict__ input, + int64_t size) { + for (int64_t d = 0; d < size; d++) { + out[d] = input[d]; + } +} + +#define PER_TENSOR 1 +#define PER_ROW 2 +#define PER_GROUP 3 + +// Store float32 GEMM result to output buffer with dequantization and dtype +// conversion. Activations are PER_ROW quantized; weights are PER_GROUP quantized. +// Uses AVX512 on capable platforms. +#if defined(CPU_CAPABILITY_AVX512) + +template +inline void store_out( + const float* y_buf, + out_dtype* c_ptr, + int64_t M, + int64_t lda, + const float* scales_a, + const float* scales_b, + const float* bias) { + float a_scale = 1.0, b_scale = 1.0; + __m512 va_scale, vb_scale; + if constexpr (act_quant_mode == PER_TENSOR) { + a_scale = *scales_a; + } + if constexpr (wei_quant_mode == PER_TENSOR) { + b_scale = *scales_b; + vb_scale = _mm512_set1_ps(b_scale); + } + for (int i = 0; i < M; ++i) { + if constexpr (act_quant_mode == PER_ROW) { + a_scale = *(scales_a + i); + } + if constexpr (act_quant_mode != PER_GROUP) { + va_scale = _mm512_set1_ps(a_scale); + } + constexpr int N_UNROLL = N / 16; + c10::ForcedUnroll{}([&](auto idx) { + constexpr int j = idx * 16; + __m512 y_vec = _mm512_loadu_ps(y_buf + i * N + j); + __m512 bias_vec = bias ? _mm512_loadu_ps(bias + j) : _mm512_setzero_ps(); + if constexpr (act_quant_mode != PER_GROUP) { + y_vec = _mm512_mul_ps(y_vec, va_scale); + } + if constexpr (wei_quant_mode == PER_ROW) { + vb_scale = _mm512_loadu_ps(scales_b + j); + } + if constexpr (wei_quant_mode != PER_GROUP) { + y_vec = _mm512_mul_ps(y_vec, vb_scale); + } + y_vec = _mm512_add_ps(y_vec, bias_vec); + if constexpr (std::is_same::value) { + _mm512_storeu_ps(c_ptr + i * lda + j, y_vec); + } else if constexpr (std::is_same::value) { + __m256i y_bf16_vec = at::vec::cvtfp32_bf16(y_vec); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c_ptr + i * lda + j), y_bf16_vec); + } else if constexpr (std::is_same::value) { + __m256i y_fp16_vec = at::vec::cvtfp32_fp16(y_vec); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c_ptr + i * lda + j), y_fp16_vec); + } else { + TORCH_CHECK(false, "Unsupported output dtype"); + } + }); + constexpr int tail_start = N / 16 * 16; + for (int j = tail_start; j < N; ++j) { + if constexpr (wei_quant_mode == PER_ROW) { + b_scale = scales_b[j]; + } + c_ptr[i * lda + j] = static_cast(y_buf[i * N + j] * a_scale * b_scale); + } + } // for M +} + +#else // no AVX512 + +template +inline void store_out( + const float* y_buf, + out_dtype* c_ptr, + int64_t M, + int64_t lda, + const float* scales_a, + const float* scales_b, + const float* bias) { + float a_scale = 1.0, b_scale = 1.0; + if constexpr (act_quant_mode == PER_TENSOR) { + a_scale = *scales_a; + } + if constexpr (wei_quant_mode == PER_TENSOR) { + b_scale = *scales_b; + } + for (int i = 0; i < M; ++i) { + if constexpr (act_quant_mode == PER_ROW) { + a_scale = *(scales_a + i); + } + for (int j = 0; j < N; ++j) { + if constexpr (wei_quant_mode == PER_ROW) { + b_scale = scales_b[j]; + } + float val = y_buf[i * N + j] * a_scale * b_scale; + if (bias) val += bias[j]; + c_ptr[i * lda + j] = static_cast(val); + } + } +} + +#endif // CPU_CAPABILITY_AVX512 + +} // anonymous namespace + +// FP8 W8A8 fused MoE kernel. +// +// hidden_states: [M, K] FP8 E4M3, pre-quantized activation +// packed_w1: [E, 2*Nc, Kc, BLOCK_K, BLOCK_N] FP8 W8A8 prepack (float8_linear_prepack_cpu) +// packed_w2: [E, K, N_packed] FP8 W8A16 VNNI pack (convert_weight_packed) +// As: [M] float32 per-token activation scales +// w1s: [E, num_groups, 2N] float32 per-group weight scales for w1 +// w2s: [E, ...] float32 per-group weight scales for w2 +// +// Stage 1: hidden_states(FP8) @ w1(FP8) → ic0(BF16) [M*topk, 2N] +// Uses FP8×FP8 GEMM (tinygemm_kernel with FP8 inputs) +// Stage 1.5: ic1 = silu(ic0[:,:N]) * ic0[:,N:] [M*topk, N] +// Stage 2: ic1(BF16) @ w2(FP8) → ic2(BF16) [M*topk, K] +// Uses BF16×FP8 GEMM (W8A16 path, tinygemm_kernel) +// Stage 3: out = weighted_sum(ic2, topk_weights) [M, K] + +template +void fused_experts_fp8_a8_kernel_impl( + scalar_t* __restrict__ output, + scalar_t* __restrict__ ic0, + scalar_t* __restrict__ ic1, + scalar_t* __restrict__ ic2, + at::Float8_e4m3fn* __restrict__ A_tmp, + scalar_t* __restrict__ B_tmp, + float* __restrict__ C_tmp, + float* __restrict__ Ukernel_tmp, + const at::Float8_e4m3fn* __restrict__ input, + const at::Float8_e4m3fn* __restrict__ packed_w1, + const at::Float8_e4m3fn* __restrict__ packed_w2, + const float* __restrict__ As, + const float* __restrict__ w1s, + const float* __restrict__ w2s, + int64_t block_size_N, + int64_t block_size_K, + const float* __restrict__ topk_weights, + const int32_t* __restrict__ sorted_ids, + const int32_t* __restrict__ expert_ids, + const int32_t* __restrict__ offsets, + int64_t M, + int64_t N, + int64_t K, + int64_t E, + int64_t topk, + int64_t num_tokens_post_pad) { + // 1. intermediate_cache1 : [M * topk, N] + // 2. intermediate_cache2 : [M * topk, K] + // 3. A_tmp : [T, BLOCK_M * K] FP8 + // 4. C_tmp : [T, 2 * BLOCK_M * BLOCK_N] float32 + // 5. intermediate_cache0 : [M * topk, 2N] + // 6. B_tmp : [T, MAX_CACHE_BLOCK_SIZE, BLOCK_N, max(K, N)] scalar_t + // 7. Ukernel_tmp: [T, 2 * BLOCK_M * BLOCK_N] float32 + // 8. (non-FP8-brgemm only) dqA [T, BLOCK_M * BLOCK_K] + dqB [T, BLOCK_K * BLOCK_N] in BF16 + // appended after Ukernel_tmp; see moe.cpp FP8_W8A8 buffer allocation + + constexpr int64_t BLOCK_M = block_size_m(); + constexpr int64_t BLOCK_N_VAL = block_size_n(); // = 32 + const int64_t KB = K / BLOCK_K; // K blocks + + int64_t B_tmp_size_per_thread = MAX_CACHE_BLOCK_SIZE * BLOCK_N_VAL * std::max(K, N); + + // stage 1: intermediate_cache0 = hidden_states @ w1 + const int64_t MB = div_up(num_tokens_post_pad, BLOCK_M); + const int64_t NB = div_up(N, BLOCK_N_VAL); + + int64_t scale_size_N = div_up(2 * N, block_size_N); + int64_t scale_size_K = div_up(K, block_size_K); + int64_t blocks_n_per_group = block_size_N / BLOCK_N_VAL; + int64_t num_groups = div_up(K, block_size_K); // G + int64_t blocks_k_per_group = block_size_K / BLOCK_K; + const int64_t stride_e = 2 * N * K; // per-expert stride for w1 (total elements) + bool use_brgemm = true; + int64_t num_thread = at::get_num_threads(); + + // For the BF16 fallback path (no native FP8 brgemm), tinygemm_kernel needs temporary + // BF16 buffers to dequantize FP8 inputs before running AMX-BF16 brgemm. + // These are placed after the Ukernel_tmp region (see moe.cpp FP8_W8A8 buffer allocation). +#ifndef CPUBLAS_BRGEMM_F8F8F32 + // dqA: [T, BLOCK_M, BLOCK_K] BF16 + // dqB: [T, BLOCK_K, BLOCK_N] BF16 + // Layout: immediately after Ukernel_tmp[num_thread * 2 * BLOCK_M * BLOCK_N_VAL] + at::BFloat16* __restrict__ dq_base = reinterpret_cast( + Ukernel_tmp + num_thread * 2 * BLOCK_M * BLOCK_N_VAL); + const int64_t dq_per_thread = BLOCK_M * BLOCK_K + BLOCK_K * BLOCK_N_VAL; +#endif + + at::parallel_for(0, MB * NB, 1, [&](int64_t begin, int64_t end) { + int tid = get_thread_num(); + // A_tmp: FP8 activation buffer for this thread [BLOCK_M, K] + at::Float8_e4m3fn* __restrict__ A = A_tmp + tid * BLOCK_M * K; + float* __restrict__ C0 = C_tmp + tid * 2 * BLOCK_M * BLOCK_N_VAL; + float* __restrict__ C1 = C0 + BLOCK_M * BLOCK_N_VAL; + alignas(64) float As_[BLOCK_M]; + float* __restrict__ ukernel_buf_1 = Ukernel_tmp + tid * 2 * BLOCK_M * BLOCK_N_VAL; +#ifndef CPUBLAS_BRGEMM_F8F8F32 + at::BFloat16* __restrict__ dqA_buf = dq_base + tid * dq_per_thread; + at::BFloat16* __restrict__ dqB_buf = dqA_buf + BLOCK_M * BLOCK_K; +#else + at::BFloat16* __restrict__ dqA_buf = nullptr; + at::BFloat16* __restrict__ dqB_buf = nullptr; +#endif + float* __restrict__ ukernel_buf_2 = ukernel_buf_1 + BLOCK_M * BLOCK_N_VAL; + auto ldsa = 1; // activation per-row quantized + + for (int64_t i = begin; i < end; ++i) { + int64_t mb = i / NB; + int64_t nb = i % NB; + int64_t nb1 = nb + NB; // gate is first N columns, up is second N columns + + int64_t n_size = std::min(N - nb * BLOCK_N_VAL, BLOCK_N_VAL); + + // B (weight) for expert expert_id: gate+up fused [2N, K] prepackaged as FP8 + int32_t expert_id = expert_ids[mb]; + const at::Float8_e4m3fn* __restrict__ B = packed_w1 + expert_id * stride_e; + const float* __restrict__ Bs = w1s + expert_id * num_groups * (2 * N); + + // Load activation tile for this M-block + const int32_t* A_ids = sorted_ids + mb * BLOCK_M; + int64_t m_size = offsets[mb + 1] - offsets[mb]; + + for (int64_t m = 0; m < m_size; ++m) { + int32_t index = A_ids[m] / topk; + copy_stub2(A + m * K, input + index * K, K); + As_[m] = As[index]; + } + + const int64_t offset = offsets[mb]; + + // Gate path (nb): w1[:N, :] + zero_buffer(C0, BLOCK_M * BLOCK_N_VAL); + for (int kci = 0; kci < KB; ++kci) { + // FP8×FP8 tinygemm: act_quant=PER_ROW(2), wei_quant=PER_GROUP(3) + tinygemm_kernel( + /* C */ C0, + /* A */ A + kci * BLOCK_K, + /* scales_a*/ As_ + kci / blocks_k_per_group, + /* B */ B + (nb * KB + kci) * BLOCK_K * BLOCK_N_VAL, + /* scales_b*/ Bs + nb * BLOCK_N_VAL * num_groups + kci / blocks_k_per_group * BLOCK_N_VAL, + /* M */ m_size, + /* K */ BLOCK_K, + /* lda */ K, + /* ldc */ BLOCK_N_VAL, + /* ldsa */ ldsa, + /* ukernel */ ukernel_buf_1, + /* dqA_buf */ dqA_buf, + /* dqB_buf */ dqB_buf); + } + store_out( + C0, + ic0 + offset * 2 * N + nb * BLOCK_N_VAL, + m_size, + 2 * N, // lda + As_, + nullptr, + nullptr); + + // Up path (nb1): w1[N:, :] + zero_buffer(C1, BLOCK_M * BLOCK_N_VAL); + for (int kci = 0; kci < KB; ++kci) { + tinygemm_kernel( + /* C */ C1, + /* A */ A + kci * BLOCK_K, + /* scales_a*/ As_ + kci / blocks_k_per_group, + /* B */ B + (nb1 * KB + kci) * BLOCK_K * BLOCK_N_VAL, + /* scales_b*/ Bs + nb1 * BLOCK_N_VAL * num_groups + kci / blocks_k_per_group * BLOCK_N_VAL, + /* M */ m_size, + /* K */ BLOCK_K, + /* lda */ K, + /* ldc */ BLOCK_N_VAL, + /* ldsa */ ldsa, + /* ukernel */ ukernel_buf_2, + /* dqA_buf */ dqA_buf, + /* dqB_buf */ dqB_buf); + } + store_out( + C1, + ic0 + offset * 2 * N + nb1 * BLOCK_N_VAL, + m_size, + 2 * N, // lda + As_, + nullptr, + nullptr); + } + if (use_brgemm) { + at::native::cpublas::brgemm_release(); + } + }); + + // stage 1.5: intermediate_cache1 = silu(ic0[:, :N]) * ic0[:, N:] + at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) { + for (int64_t m = begin; m < end; ++m) { + silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N); + } + }); + + // stage 2: intermediate_cache2 = intermediate_cache1 @ w2 + // w2: [E, K, N] FP8, VNNI packed; treated as W8A16 (BF16 act × FP8 weight) + const int64_t OC = K; // output channels (hidden dim) + const int64_t IC = N; // input channels (expert intermediate dim) + const int64_t MB2 = MB; + const int64_t NB2 = div_up(OC, BLOCK_N_VAL); + scale_size_N = div_up(K, block_size_N); + scale_size_K = div_up(N, block_size_K); + const int64_t stride_e2 = OC * IC; // per-expert stride for w2 + const int64_t stride_oc = IC; + + int64_t avg_M = std::max(int64_t(1), M * topk / E); + bool use_brgemm2 = can_use_brgemm(avg_M); + + parallel_2d(MB2, NB2, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) { + int tid = get_thread_num(); + alignas(64) scalar_t C[BLOCK_M * BLOCK_K]; + + loop_2d(mb0, mb1, nb0, nb1, BLOCK_N_VAL * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) { + int64_t m_size = offsets[mb + 1] - offsets[mb]; + int64_t n_size = std::min(OC - nb * BLOCK_N_VAL, BLOCK_N_VAL); + + // A from ic1: [M*topk, N] in sorted order (contiguous after stage 1.5) + const scalar_t* __restrict__ A2 = ic1 + offsets[mb] * N; + const int32_t* A_ids = sorted_ids + mb * BLOCK_M; + + // B from w2: [E, K, N] VNNI packed, offset to expert + output block + int32_t expert_id = expert_ids[mb]; + const at::Float8_e4m3fn* __restrict__ B2 = + packed_w2 + expert_id * stride_e2 + nb * BLOCK_N_VAL * stride_oc; + const float* __restrict__ Bs2 = + w2s + expert_id * scale_size_N * scale_size_K + (nb / blocks_n_per_group) * scale_size_K; + + // Only unpack B on first visit or new expert + int32_t pre_expert_id = mb == 0 ? -1 : expert_ids[mb - 1]; + bool do_unpack = (mb == mb0) || (expert_id != pre_expert_id); + + tinygemm_kernel( + /* A */ A2, + /* B */ B2, + /* C */ C, + /* Btmp */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N_VAL * IC, + /* Ctmp */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N_VAL, + /* Bbias */ nullptr, + /* scale */ Bs2, + /* M */ m_size, + /* N */ n_size, + /* K */ IC, + /* lda */ IC, + /* ldb */ n_size, + /* ldc */ BLOCK_N_VAL, + /* brg */ use_brgemm2, + /* block_sz_K */ block_size_K, + /* do_unpack */ do_unpack); + + // Scatter ic2 in original token order, weighting by topk_weights + for (int64_t m = 0; m < m_size; ++m) { + int32_t index = A_ids[m]; + float weight = topk_weights[index]; + copy_mul_stub(ic2 + index * K + nb * BLOCK_N_VAL, C + m * BLOCK_N_VAL, weight, n_size); + } + }); + + if (use_brgemm2) { + at::native::cpublas::brgemm_release(); + } + }); + + // stage 3: output = sum(ic2, dim=topk) [M, K] + at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) { + for (int64_t m = begin; m < end; ++m) { + sum_stub(output + m * K, ic2 + m * topk * K, topk, K); + } + }); +} + +#define INSTANTIATE_MOE_FP8_A8_TEMPLATE(TYPE) \ + template void fused_experts_fp8_a8_kernel_impl( \ + TYPE* __restrict__ output, \ + TYPE* __restrict__ ic0, \ + TYPE* __restrict__ ic1, \ + TYPE* __restrict__ ic2, \ + at::Float8_e4m3fn* __restrict__ A_tmp, \ + TYPE* __restrict__ B_tmp, \ + float* __restrict__ C_tmp, \ + float* __restrict__ Ukernel_tmp, \ + const at::Float8_e4m3fn* __restrict__ input, \ + const at::Float8_e4m3fn* __restrict__ packed_w1, \ + const at::Float8_e4m3fn* __restrict__ packed_w2, \ + const float* __restrict__ As, \ + const float* __restrict__ w1s, \ + const float* __restrict__ w2s, \ + int64_t block_size_N, \ + int64_t block_size_K, \ + const float* __restrict__ topk_weights, \ + const int32_t* __restrict__ sorted_ids, \ + const int32_t* __restrict__ expert_ids, \ + const int32_t* __restrict__ offsets, \ + int64_t M, \ + int64_t N, \ + int64_t K, \ + int64_t E, \ + int64_t topk, \ + int64_t num_tokens_post_pad) + +INSTANTIATE_MOE_FP8_A8_TEMPLATE(at::BFloat16); +INSTANTIATE_MOE_FP8_A8_TEMPLATE(at::Half); diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index 88170068de..accd5fcec3 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -73,6 +73,7 @@ void fused_experts_cpu(at::Tensor& out, at::Tensor& hidden_states, const std::optional& w2_scale, const std::optional& w1_zero, const std::optional& w2_zero, + const std::optional& a1_scale, const std::optional> block_size, const std::optional& w1_bias, const std::optional& w2_bias, @@ -91,6 +92,18 @@ at::Tensor fp8_scaled_mm_cpu(at::Tensor& mat1, at::Tensor& mat2, const std::optional& bias, at::ScalarType out_dtype, bool is_vnni); +// Adapted from sglang: FP8 W8A8 kernels +std::tuple float8_linear_prepack_impl( + const at::Tensor& weight, const at::Tensor& scales); +at::Tensor fp8_scaled_mm_with_quant(const at::Tensor& act, + const std::optional& act_scales, + bool channelwise, const at::Tensor& weight, + const at::Tensor& weight_scales, + const std::optional& bias, + at::ScalarType output_dtype); +std::tuple quantize_fp8e4m3_vec( + const at::Tensor& t, bool channelwise, c10::optional scale_opt); + // Adapted from sglang: MLA CPU kernels (AMX-only) void decode_attention_cpu(at::Tensor& query, at::Tensor& k_buffer, at::Tensor& v_buffer, at::Tensor& output, @@ -281,6 +294,7 @@ void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input, const std::string& activation); bool cpu_attn_has_isa(const std::string& isa); +bool cpu_has_amx_fp8(); torch::Tensor get_scheduler_metadata( const int64_t num_req, const int64_t num_heads_q, @@ -611,7 +625,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "fused_experts_cpu(Tensor(a0!) out, Tensor hidden_states, Tensor w1, " "Tensor w2, Tensor topk_weights, Tensor topk_ids, " "int moe_comp_method, Tensor? w1_scale, Tensor? w2_scale, " - "Tensor? w1_zero, Tensor? w2_zero, int[]? block_size, " + "Tensor? w1_zero, Tensor? w2_zero, Tensor? a1_scale, int[]? block_size, " "Tensor? w1_bias, Tensor? w2_bias, float? alpha, float? limit, " "bool is_vnni) -> " "()"); @@ -629,7 +643,23 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "bool is_vnni) -> Tensor"); ops.impl("fp8_scaled_mm_cpu", torch::kCPU, &fp8_scaled_mm_cpu); - // Adapted from sglang: casual_conv1d kernels + // Adapted from sglang: FP8 W8A8 kernels + ops.def( + "float8_linear_prepack_cpu(Tensor weight, Tensor scales) -> (Tensor, " + "Tensor)"); + ops.impl("float8_linear_prepack_cpu", torch::kCPU, + &float8_linear_prepack_impl); + ops.def( + "fp8_scaled_mm_with_quant(Tensor act, Tensor? act_scales, bool " + "channelwise, " + "Tensor weight, Tensor weight_scales, Tensor? bias, ScalarType " + "output_dtype) -> Tensor"); + ops.impl("fp8_scaled_mm_with_quant", torch::kCPU, &fp8_scaled_mm_with_quant); + ops.def( + "quantize_fp8e4m3_vec(Tensor input, bool channelwise, Tensor? scale_opt) " + "-> (Tensor, Tensor)"); + ops.impl("quantize_fp8e4m3_vec", torch::kCPU, &quantize_fp8e4m3_vec); + ops.def("causal_conv1d_weight_pack(Tensor weight) -> Tensor"); ops.impl("causal_conv1d_weight_pack", torch::kCPU, &causal_conv1d_weight_pack); @@ -832,6 +862,9 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "-> (Tensor, Tensor)"); ops.impl("fused_gdn_gating_cpu", torch::kCPU, &fused_gdn_gating_cpu); + // Pure runtime CPU ISA capability check. + ops.def("cpu_has_amx_fp8() -> bool", &cpu_has_amx_fp8); + // CPU attention kernels ops.def("cpu_attn_has_isa(str isa) -> bool", &cpu_attn_has_isa); ops.def( diff --git a/tests/kernels/moe/test_cpu_quant_fused_moe.py b/tests/kernels/moe/test_cpu_quant_fused_moe.py index 99cae4534b..e89fbb55a0 100644 --- a/tests/kernels/moe/test_cpu_quant_fused_moe.py +++ b/tests/kernels/moe/test_cpu_quant_fused_moe.py @@ -1114,5 +1114,203 @@ def test_int8_w8a8_cpu_fused_moe(M, N, K, E, topk, seed, is_vnni, inplace): ) +# =========================================================================== +# FP8 W8A8 MoE +# =========================================================================== + + +requires_cpu_fp8_w8a8 = pytest.mark.skipif( + not ops.cpu_has_amx_fp8(), + reason="requires native AMX-FP8 hardware", +) + +FP8_W8A8_MAX = torch.finfo(torch.float8_e4m3fn).max +FP8_W8A8_QUANT_GROUP = 128 # quantization group size for both K and N dimensions + + +def _make_fp8_w8a8_weight_w13( + E: int, two_n: int, K: int, group_K: int = FP8_W8A8_QUANT_GROUP +) -> tuple[torch.Tensor, torch.Tensor]: + """Create w13 [E, 2N, K] FP8 with group scales [E, 2N, G] where G = K//group_K.""" + G = K // group_K + w_list, s_list = [], [] + for _ in range(E): + w_f32 = torch.randn(two_n, K) + # Per-row, per-K-group quantization: scale [2N, G] + w_re = w_f32.view(two_n, G, group_K) + abs_max = w_re.abs().amax(dim=2, keepdim=True).clamp(min=1e-7) # [2N, G, 1] + scale = (abs_max / FP8_W8A8_MAX).squeeze(2) # [2N, G] + w_q = ( + (w_re / abs_max).clamp(-FP8_W8A8_MAX, FP8_W8A8_MAX).to(torch.float8_e4m3fn) + ) + w_list.append(w_q.view(two_n, K).contiguous()) + s_list.append(scale.float()) + return torch.stack(w_list), torch.stack(s_list) # [E, 2N, K], [E, 2N, G] + + +def _make_fp8_w8a8_weight_w2( + E: int, + K: int, + N: int, + group_K: int = FP8_W8A8_QUANT_GROUP, + group_N: int = FP8_W8A8_QUANT_GROUP, +) -> tuple[torch.Tensor, torch.Tensor]: + """Create w2 [E, K, N] FP8 with block scales [E, K//gK, N//gN].""" + nK = K // group_K + nN = N // group_N + w_list, s_list = [], [] + for _ in range(E): + w_f32 = torch.randn(K, N) + w_re = w_f32.view(nK, group_K, nN, group_N) + abs_max = w_re.abs().amax(dim=(1, 3), keepdim=True).clamp(min=1e-7) + scale = (abs_max / FP8_W8A8_MAX).squeeze(1).squeeze(2) # [nK, nN] + w_q = ( + (w_re / abs_max).clamp(-FP8_W8A8_MAX, FP8_W8A8_MAX).to(torch.float8_e4m3fn) + ) + w_list.append(w_q.view(K, N).contiguous()) + s_list.append(scale.float()) + return torch.stack(w_list), torch.stack(s_list) # [E, K, N], [E, K//gK, N//gN] + + +def _pack_fp8_w8a8_w13( + w13: torch.Tensor, # [E, 2N, K] FP8 + w13_scale: torch.Tensor, # [E, 2N, G] float32 +) -> tuple[torch.Tensor, torch.Tensor]: + """Pack w13 for CPU FP8 W8A8 MoE kernel using float8_linear_prepack_cpu.""" + E = w13.size(0) + packed_list, scale_list = [], [] + for i in range(E): + pw, ps = torch.ops._C.float8_linear_prepack_cpu( + w13[i].contiguous(), w13_scale[i].contiguous() + ) + packed_list.append(pw) + scale_list.append(ps) + return torch.stack(packed_list), torch.stack(scale_list) + + +def _ref_fp8_w8a8_moe( + hidden_states: torch.Tensor, # BF16 [M, K_hidden] (after FP8 round-trip dequant) + w13: torch.Tensor, # FP8 [E, 2N, K_hidden] + w13_scale: torch.Tensor, # float32 [E, 2N, G] (G = K_hidden//group_K) + w2: torch.Tensor, # FP8 [E, K_hidden, N] + w2_scale: torch.Tensor, # float32 [E, K_hidden//gK, N//gN] + topk_weights: torch.Tensor, # float32 [M, top_k] + topk_ids: torch.Tensor, # int64 [M, top_k] + N: int, + group_K: int = FP8_W8A8_QUANT_GROUP, + group_N: int = FP8_W8A8_QUANT_GROUP, +) -> torch.Tensor: + """Reference implementation: dequantize weights then compute MoE.""" + M, K_hidden = hidden_states.shape + top_k = topk_ids.shape[1] + output = torch.zeros(M, K_hidden, dtype=torch.bfloat16) + + for tok in range(M): + x = hidden_states[tok].float() # [K_hidden] + for k_idx in range(top_k): + expert_id = topk_ids[tok, k_idx].item() + w = topk_weights[tok, k_idx].item() + + # Stage 1: x @ w13.T — gate and up projections + w13_e = w13[expert_id].float() # [2N, K_hidden] + ws_e = w13_scale[expert_id].float() # [2N, G] + G = ws_e.shape[1] + w13_dq = torch.zeros_like(w13_e) + for g in range(G): + c0, c1 = g * group_K, (g + 1) * group_K + w13_dq[:, c0:c1] = w13_e[:, c0:c1] * ws_e[:, g : g + 1] + gate_up = (x.unsqueeze(0) @ w13_dq.T).squeeze(0) # [2N] + + gate = gate_up[:N] + up = gate_up[N:] + act = F.silu(gate) * up # [N] + + # Stage 2: act @ w2.T — down projection + w2_e = w2[expert_id].float() # [K_hidden, N] + ws2_e = w2_scale[expert_id].float() # [K_hidden//gK, N//gN] + nK2 = K_hidden // group_K + nN2 = N // group_N + w2_dq = torch.zeros_like(w2_e) + for gi in range(nK2): + for gj in range(nN2): + r0, r1 = gi * group_K, (gi + 1) * group_K + c0, c1 = gj * group_N, (gj + 1) * group_N + w2_dq[r0:r1, c0:c1] = w2_e[r0:r1, c0:c1] * ws2_e[gi, gj] + out_e = (act.unsqueeze(0) @ w2_dq.T).squeeze(0) # [K_hidden] + + output[tok] += (w * out_e).to(torch.bfloat16) + + return output + + +@requires_cpu_fp8_w8a8 +@pytest.mark.parametrize( + "E,N,K,M,top_k", + [ + # N and K must be multiples of FP8_W8A8_QUANT_GROUP=128 + (4, 128, 256, 1, 2), # single token (decode phase) + (4, 128, 256, 8, 2), + (8, 256, 512, 16, 2), + (4, 256, 512, 4, 1), + ], +) +def test_fp8_w8a8_cpu_fused_moe(E: int, N: int, K: int, M: int, top_k: int): + """Test fused_experts_cpu FP8_W8A8 shape/dtype and accuracy vs a + dequantized BF16 reference.""" + set_random_seed(42) + + w13, w13_scale = _make_fp8_w8a8_weight_w13(E, 2 * N, K) + w2, w2_scale = _make_fp8_w8a8_weight_w2(E, K, N) + + packed_w13, packed_w13_scale = _pack_fp8_w8a8_w13(w13, w13_scale) + packed_w2 = _prepack_experts(w2) + + hidden_states = torch.randn(M, K, dtype=torch.bfloat16) * 0.1 + x_fp8, x_scales = torch.ops._C.quantize_fp8e4m3_vec(hidden_states, True, None) + assert x_fp8.dtype == torch.float8_e4m3fn + assert x_scales.shape == (M,) + + gen = torch.Generator().manual_seed(0) + topk_weights = torch.softmax(torch.randn(M, top_k, generator=gen), dim=-1).float() + topk_ids = torch.zeros(M, top_k, dtype=torch.int32) + for i in range(M): + topk_ids[i] = torch.randperm(E, generator=torch.Generator().manual_seed(i))[ + :top_k + ].int() + + output = torch.empty(M, K, dtype=torch.bfloat16) + ops.fused_experts_cpu( + output, + x_fp8, + packed_w13, + packed_w2, + topk_weights, + topk_ids, + ops.CPUQuantMethod.FP8_W8A8, + packed_w13_scale, # w1_scale: [E, Nc, G, BLOCK_N] + w2_scale, # w2_scale: [E, K//128, N//128] + None, # w1_zero + None, # w2_zero + [FP8_W8A8_QUANT_GROUP, FP8_W8A8_QUANT_GROUP], # block_size + None, # w1_bias + None, # w2_bias + None, # alpha + None, # limit + True, # is_vnni (weights already packed) + x_scales, # a1_scale: per-token FP8 activation scales + ) + + assert output.shape == (M, K) + assert output.dtype == torch.bfloat16 + assert not torch.isnan(output).any() + assert not torch.isinf(output).any() + + x_dequant = (x_fp8.float() * x_scales.unsqueeze(1)).bfloat16() + ref = _ref_fp8_w8a8_moe( + x_dequant, w13, w13_scale, w2, w2_scale, topk_weights, topk_ids.long(), N + ) + torch.testing.assert_close(output.float(), ref.float(), atol=1e-2, rtol=1e-2) + + if __name__ == "__main__": sys.exit(pytest.main([__file__, "-v"])) diff --git a/tests/kernels/quantization/test_cpu_fp8_scaled_mm.py b/tests/kernels/quantization/test_cpu_fp8_scaled_mm.py index e4dadb53bd..0b0d12a66c 100644 --- a/tests/kernels/quantization/test_cpu_fp8_scaled_mm.py +++ b/tests/kernels/quantization/test_cpu_fp8_scaled_mm.py @@ -1,6 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for CPU FP8 W8A16 block-scaled GEMM kernel (fp8_scaled_mm_cpu). +"""Tests for CPU FP8 scaled-mm GEMM kernels. + +Covers: + * fp8_scaled_mm_cpu (W8A16) + * fp8_scaled_mm_with_quant (W8A8) Run `pytest tests/kernels/quantization/test_cpu_fp8_scaled_mm.py -v`. """ @@ -23,19 +27,31 @@ from vllm.platforms import current_platform if not current_platform.is_cpu(): pytest.skip("skipping CPU-only tests", allow_module_level=True) -if not ops._supports_cpu_fp8_w8a16: - pytest.skip("fp8_scaled_mm_cpu op not available", allow_module_level=True) +requires_cpu_fp8_w8a16 = pytest.mark.skipif( + not ops._supports_cpu_fp8_w8a16, reason="fp8_scaled_mm_cpu op not available" +) +requires_cpu_fp8_w8a8 = pytest.mark.skipif( + not ops.cpu_has_amx_fp8(), + reason="requires native AMX-FP8 hardware", +) if not torch.cpu._is_amx_tile_supported(): pytest.skip("requires AMX tile support", allow_module_level=True) +FP8_MAX = torch.finfo(torch.float8_e4m3fn).max BLOCK_SIZE = [128, 128] +GROUP_SIZE = 128 def cdiv(a: int, b: int) -> int: return -(a // -b) +# --------------------------------------------------------------------------- +# W8A16 helpers +# --------------------------------------------------------------------------- + + def quantize_weight_block_fp8( weight: torch.Tensor, block_size: list[int], @@ -126,7 +142,7 @@ def ref_fp8_block_scaled_mm( # --------------------------------------------------------------------------- -# Test parameters +# W8A16 test: fp8_scaled_mm_cpu # --------------------------------------------------------------------------- M_SIZES = [1, 4, 16, 64, 128] # (N, K) — weight shape is [N, K], output has N columns. @@ -141,6 +157,7 @@ NK_SIZES = [ ] +@requires_cpu_fp8_w8a16 @pytest.mark.parametrize("M", M_SIZES) @pytest.mark.parametrize("N,K", NK_SIZES) @pytest.mark.parametrize("use_bias", [False, True]) @@ -200,6 +217,77 @@ def ref_fp8_per_tensor_scaled_mm( return out.to(out_dtype) +# --------------------------------------------------------------------------- +# W8A8 helpers +# --------------------------------------------------------------------------- + + +def quantize_weight_per_tensor( + weight: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize weight [N, K] to FP8 with a single per-tensor scale, broadcast + to [N, 1] (the shape float8_linear_prepack_cpu expects).""" + N = weight.shape[0] + abs_max = weight.abs().max() + scale = (abs_max / FP8_MAX).clamp(min=1e-7) + q = (weight.float() / scale).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn) + return q, scale.expand(N).contiguous().view(N, 1) + + +def quantize_weight_per_channel( + weight: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize weight [N, K] to FP8 with per-channel (per-row) scales [N, 1].""" + abs_max = weight.abs().amax(dim=1, keepdim=True) + scale = (abs_max / FP8_MAX).clamp(min=1e-7) + q = (weight.float() / scale).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn) + return q, scale.float() + + +def quantize_weight_per_group( + weight: torch.Tensor, + group_size: int = GROUP_SIZE, +) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize weight [N, K] to FP8 with per-group (128-block) scales [N, G].""" + N, K = weight.shape + G = K // group_size + w_groups = weight.view(N, G, group_size) + abs_max = w_groups.abs().amax(dim=2, keepdim=True) + scale = (abs_max / FP8_MAX).clamp(min=1e-7) + q = (w_groups.float() / scale).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn) + return q.view(N, K), scale.squeeze(2).float() + + +def quantize_act_per_token(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Reference per-token (row) dynamic FP8 quantization of activations.""" + x_f = x.float() + scale = (x_f.abs().amax(dim=1, keepdim=True) / FP8_MAX).clamp(min=1e-7) + scaled = (x_f / scale).clamp(-FP8_MAX, FP8_MAX) + # Round via fp16 to match the AVX10.2 kernel path. + q = scaled.half().to(torch.float8_e4m3fn) + return q, scale + + +def reference_static_fp8_linear( + x: torch.Tensor, # BF16 [M, K] + weight_fp8: torch.Tensor, # FP8 [N, K] + weight_scale: torch.Tensor, # float32 [N, 1] + act_scale: torch.Tensor, # float32 scalar + bias: torch.Tensor | None, + out_dtype: torch.dtype, +) -> torch.Tensor: + """Static per-tensor: quantize/dequant x with act_scale, dequant weight, matmul.""" + x_scaled = (x.float() / act_scale.item()).clamp(-FP8_MAX, FP8_MAX) + # Round via fp16 to match the AVX10.2 kernel path. + x_q = x_scaled.half().to(torch.float8_e4m3fn) + x_dq = x_q.float() * act_scale.item() + w_dq = weight_fp8.float() * weight_scale.view(-1, 1) + out = torch.mm(x_dq, w_dq.t()) + if bias is not None: + out = out + bias.float() + return out.to(out_dtype) + + NK_SIZES_PER_TENSOR = [ (32, 64), (5120, 5120), @@ -274,3 +362,136 @@ def test_cpu_fp8_per_tensor_kernel_rejects_non_multiple_of_32_n(n: int): ) supported, _ = CPUFp8PerTensorScaledMMLinearKernel.can_implement(config) assert not supported + + +def reference_dynamic_fp8_linear( + x: torch.Tensor, # BF16 [M, K] + weight_fp8: torch.Tensor, # FP8 [N, K] + weight_scale: torch.Tensor, # float32 [N, G] + bias: torch.Tensor | None, + out_dtype: torch.dtype, +) -> torch.Tensor: + """Dynamic per-token act x per-group weight: quantize x per-token, + dequant weight per-group, matmul.""" + N, K = weight_fp8.shape + G = weight_scale.shape[1] + group_size = K // G + x_q, x_scale = quantize_act_per_token(x) + x_dq = x_q.float() * x_scale + w_dq = ( + weight_fp8.float().view(N, G, group_size) * weight_scale.view(N, G, 1) + ).view(N, K) + out = torch.mm(x_dq, w_dq.t()) + if bias is not None: + out = out + bias.float() + return out.to(out_dtype) + + +# --------------------------------------------------------------------------- +# W8A8 tests: fp8_scaled_mm_with_quant +# --------------------------------------------------------------------------- +W8A8_M_SIZES = [1, 4, 16, 64] +W8A8_NK_SIZES = [(128, 256), (256, 512), (64, 128)] + + +@requires_cpu_fp8_w8a8 +@pytest.mark.parametrize("static_activation", [True, False]) +@pytest.mark.parametrize("M", W8A8_M_SIZES) +@pytest.mark.parametrize("N,K", W8A8_NK_SIZES) +@pytest.mark.parametrize("use_bias", [False, True]) +def test_cpu_fp8_w8a8_per_tensor( + M: int, N: int, K: int, use_bias: bool, static_activation: bool +): + """CPUFP8W8A8ScaledMMLinearKernel with per-tensor weight scale.""" + torch.manual_seed(42) + out_dtype = torch.bfloat16 + + x = torch.randn(M, K, dtype=torch.bfloat16) + weight = torch.randn(N, K, dtype=torch.bfloat16) + bias = torch.randn(N, dtype=torch.float32) if use_bias else None + + w_fp8, w_scale = quantize_weight_per_tensor(weight) + packed_weight, packed_scale = torch.ops._C.float8_linear_prepack_cpu(w_fp8, w_scale) + + if static_activation: + act_scale = (x.float().abs().max() / FP8_MAX).clamp(min=1e-7).view(1) + kernel_out = ops.fp8_scaled_mm_with_quant( + x, act_scale, False, packed_weight, packed_scale, bias, out_dtype + ) + ref_out = reference_static_fp8_linear( + x, w_fp8, w_scale, act_scale, bias, out_dtype + ) + atol = 0.05 + else: + kernel_out = ops.fp8_scaled_mm_with_quant( + x, None, True, packed_weight, packed_scale, bias, out_dtype + ) + ref_out = reference_dynamic_fp8_linear(x, w_fp8, w_scale, bias, out_dtype) + atol = 0.05 + + assert kernel_out.dtype == out_dtype + torch.testing.assert_close(kernel_out, ref_out, rtol=0.05, atol=atol) + + +@requires_cpu_fp8_w8a8 +@pytest.mark.parametrize("static_activation", [True, False]) +@pytest.mark.parametrize("M", W8A8_M_SIZES) +@pytest.mark.parametrize("N,K", W8A8_NK_SIZES) +@pytest.mark.parametrize("use_bias", [False, True]) +def test_cpu_fp8_w8a8_per_channel( + M: int, N: int, K: int, use_bias: bool, static_activation: bool +): + """CPUFP8W8A8ScaledMMLinearKernel with per-channel weight scale.""" + torch.manual_seed(42) + out_dtype = torch.bfloat16 + + x = torch.randn(M, K, dtype=torch.bfloat16) + weight = torch.randn(N, K, dtype=torch.bfloat16) + bias = torch.randn(N, dtype=torch.float32) if use_bias else None + + w_fp8, w_scale = quantize_weight_per_channel(weight) + packed_weight, packed_scale = torch.ops._C.float8_linear_prepack_cpu(w_fp8, w_scale) + + if static_activation: + act_scale = (x.float().abs().max() / FP8_MAX).clamp(min=1e-7).view(1) + kernel_out = ops.fp8_scaled_mm_with_quant( + x, act_scale, False, packed_weight, packed_scale, bias, out_dtype + ) + ref_out = reference_static_fp8_linear( + x, w_fp8, w_scale, act_scale, bias, out_dtype + ) + atol = 0.05 + else: + kernel_out = ops.fp8_scaled_mm_with_quant( + x, None, True, packed_weight, packed_scale, bias, out_dtype + ) + ref_out = reference_dynamic_fp8_linear(x, w_fp8, w_scale, bias, out_dtype) + atol = 0.05 + + assert kernel_out.dtype == out_dtype + torch.testing.assert_close(kernel_out, ref_out, rtol=0.05, atol=atol) + + +@requires_cpu_fp8_w8a8 +@pytest.mark.parametrize("M", W8A8_M_SIZES) +@pytest.mark.parametrize("N,K", W8A8_NK_SIZES) +@pytest.mark.parametrize("use_bias", [False, True]) +def test_cpu_fp8_w8a8_per_group(M: int, N: int, K: int, use_bias: bool): + """CPUFp8W8A8BlockScaledMMKernel: per-group (128-block) weight scale + + dynamic per-token activation scale.""" + torch.manual_seed(42) + out_dtype = torch.bfloat16 + + x = torch.randn(M, K, dtype=torch.bfloat16) + weight = torch.randn(N, K, dtype=torch.bfloat16) + bias = torch.randn(N, dtype=torch.float32) if use_bias else None + + w_fp8, w_scale = quantize_weight_per_group(weight) + packed_weight, packed_scale = torch.ops._C.float8_linear_prepack_cpu(w_fp8, w_scale) + kernel_out = ops.fp8_scaled_mm_with_quant( + x, None, True, packed_weight, packed_scale, bias, out_dtype + ) + ref_out = reference_dynamic_fp8_linear(x, w_fp8, w_scale, bias, out_dtype) + + assert kernel_out.dtype == out_dtype + torch.testing.assert_close(kernel_out, ref_out, rtol=0.05, atol=0.05) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 96f47e2069..2a1654231f 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from enum import IntEnum +from functools import cache from typing import TYPE_CHECKING, Literal import torch @@ -3384,6 +3385,7 @@ class CPUQuantMethod(IntEnum): FP8_W8A16 = 2 INT4_W4A8 = 3 MXFP4 = 4 + FP8_W8A8 = 5 if hasattr(torch.ops._C, "dynamic_4bit_int_moe"): @@ -3422,6 +3424,7 @@ def fused_experts_cpu( alpha: float | None = None, limit: float | None = None, is_vnni: bool = True, + a1_scale: torch.Tensor | None = None, ) -> None: torch.ops._C.fused_experts_cpu( out, @@ -3435,6 +3438,7 @@ def fused_experts_cpu( w2_scale, w1_zero, w2_zero, + a1_scale, block_size, w1_bias, w2_bias, @@ -3564,6 +3568,46 @@ def fp8_scaled_mm_cpu( ) +# FP8 W8A8 CPU kernels +@cache +def cpu_has_amx_fp8() -> bool: + """Whether this CPU has native AMX-FP8 MMA support.""" + if not hasattr(torch.ops._C, "cpu_has_amx_fp8"): + return False + return bool(torch.ops._C.cpu_has_amx_fp8()) + + +if hasattr(torch.ops._C, "fp8_scaled_mm_with_quant"): + + @register_fake("_C::fp8_scaled_mm_with_quant") + def fp8_scaled_mm_with_quant_fake( + act: torch.Tensor, + act_scales: torch.Tensor | None, + channelwise: bool, + weight: torch.Tensor, + weight_scales: torch.Tensor, + bias: torch.Tensor | None, + output_dtype: torch.dtype, + ) -> torch.Tensor: + M = act.reshape(-1, act.size(-1)).size(0) + N = weight.size(0) * weight.size(-1) + return torch.empty((M, N), dtype=output_dtype, device=act.device) + + +def fp8_scaled_mm_with_quant( + act: torch.Tensor, + act_scales: torch.Tensor | None, + channelwise: bool, + weight: torch.Tensor, + weight_scales: torch.Tensor, + bias: torch.Tensor | None, + output_dtype: torch.dtype, +) -> torch.Tensor: + return torch.ops._C.fp8_scaled_mm_with_quant( + act, act_scales, channelwise, weight, weight_scales, bias, output_dtype + ) + + def chunk_gated_delta_rule_cpu( query: torch.Tensor, key: torch.Tensor, diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 8c57059e57..fb7a170b82 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -186,6 +186,8 @@ from vllm.model_executor.kernels.linear.scaled_mm.b12x import ( from vllm.model_executor.kernels.linear.scaled_mm.cpu import ( CPUFp8BlockScaledMMKernel, CPUFp8PerTensorScaledMMLinearKernel, + CPUFp8W8A8BlockScaledMMKernel, + CPUFP8W8A8ScaledMMLinearKernel, CPUInt8ScaledMMLinearKernel, ) from vllm.model_executor.kernels.linear.scaled_mm.cutlass import ( @@ -436,6 +438,7 @@ _POSSIBLE_FP8_KERNELS: dict[PlatformEnum, list[type[FP8ScaledMMLinearKernel]]] = ChannelWiseTorchFP8ScaledMMLinearKernel, ], PlatformEnum.CPU: [ + CPUFP8W8A8ScaledMMLinearKernel, CPUFp8PerTensorScaledMMLinearKernel, PerTensorTorchFP8ScaledMMLinearKernel, ChannelWiseTorchFP8ScaledMMLinearKernel, @@ -468,6 +471,7 @@ _POSSIBLE_FP8_BLOCK_KERNELS: dict[ TritonFp8BlockScaledMMKernel, ], PlatformEnum.CPU: [ + CPUFp8W8A8BlockScaledMMKernel, # W8A8 preferred; falls back to W8A16 below CPUFp8BlockScaledMMKernel, ], PlatformEnum.XPU: [ diff --git a/vllm/model_executor/kernels/linear/scaled_mm/__init__.py b/vllm/model_executor/kernels/linear/scaled_mm/__init__.py index f4a062b7f2..75db4bcaa9 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/__init__.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/__init__.py @@ -10,6 +10,7 @@ from vllm.model_executor.kernels.linear.scaled_mm.BlockScaledMMLinearKernel impo from vllm.model_executor.kernels.linear.scaled_mm.cpu import ( CPUFp8BlockScaledMMKernel, CPUFp8PerTensorScaledMMLinearKernel, + CPUFP8W8A8ScaledMMLinearKernel, CPUInt8ScaledMMLinearKernel, ) from vllm.model_executor.kernels.linear.scaled_mm.cutlass import ( @@ -69,5 +70,6 @@ __all__ = [ "Fp8BlockScaledMMLinearKernel", "CPUFp8BlockScaledMMKernel", "CPUFp8PerTensorScaledMMLinearKernel", + "CPUFP8W8A8ScaledMMLinearKernel", "XPUFp8BlockScaledMMKernel", ] diff --git a/vllm/model_executor/kernels/linear/scaled_mm/cpu.py b/vllm/model_executor/kernels/linear/scaled_mm/cpu.py index e72dffcbbc..1d2ede1853 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/cpu.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/cpu.py @@ -270,12 +270,17 @@ class CPUFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): # Skip the base class process (FP8 padding / fnuz normalization) # which is GPU-oriented. Instead, VNNI-prepack weights for AMX. params = self._get_layer_params(layer) - packed_weight = torch.ops._C.convert_weight_packed(params.weight) - replace_parameter( - layer, - params.WEIGHT, - torch.nn.Parameter(packed_weight, requires_grad=False), - ) + + # MLA's kv_b_proj (see `_cpu_skip_gemm_dispatch`): MLA reads this + # weight to derive W_UK/W_UV, so keep the checkpoint [N, K] layout and + # let the kernel pack on the fly instead. + if not getattr(layer, "_cpu_skip_gemm_dispatch", False): + packed_weight = torch.ops._C.convert_weight_packed(params.weight) + replace_parameter( + layer, + params.WEIGHT, + torch.nn.Parameter(packed_weight, requires_grad=False), + ) # Re-wrap scale as a plain Parameter so the kernel can read it # without weight-loader metadata interfering. @@ -320,7 +325,7 @@ class CPUFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): list(self.weight_group_shape), bias, x.dtype, - True, # is_vnni (weight already prepacked) + not getattr(layer, "_cpu_skip_gemm_dispatch", False), # is_vnni ) return out.reshape(x.shape[:-1] + (out.size(-1),)) if x.dim() > 2 else out @@ -430,3 +435,294 @@ class CPUFp8PerTensorScaledMMLinearKernel(FP8ScaledMMLinearKernel): raise NotImplementedError( "CPUFp8PerTensorScaledMMLinearKernel overrides apply_weights directly." ) + + +class CPUFp8W8A8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): + """FP8 W8A8 block-quantized GEMM with dynamic per-token activation quant on CPU.""" + + # Activation quant is done inside fp8_scaled_mm_with_quant; + # the base class must not apply a second quant step. + apply_input_quant = False + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.is_cpu(): + return False, "requires CPU platform." + if not torch.cpu._is_amx_tile_supported(): + return False, "requires AMX tile support (Sapphire Rapids or newer)." + if not ops.cpu_has_amx_fp8(): + return False, ( + "requires native AMX-FP8 hardware; falls back " + "to CPUFp8BlockScaledMMKernel (W8A16) otherwise." + ) + return True, None + + @classmethod + def can_implement( + cls, config: FP8ScaledMMLinearLayerConfig + ) -> tuple[bool, str | None]: + # Same block-shape requirements as the W8A16 sibling kernel. + weight_gs = config.weight_quant_key.scale.group_shape + if weight_gs.col <= 0 or weight_gs.col != 128: + return False, ( + "CPU FP8 W8A8 block kernel requires K-dimension block size of 128, " + f"got {weight_gs.col}." + ) + if weight_gs.row <= 0 or weight_gs.row % 32 != 0: + return False, ( + "CPU FP8 W8A8 block kernel requires N-dimension block size to be " + f"a positive multiple of 32, got {weight_gs.row}." + ) + if config.out_dtype not in (torch.bfloat16, torch.float32): + return False, "Only bfloat16/float32 output dtype supported." + n = config.weight_shape[0] + if n % 32 != 0: + # float8_linear_prepack_cpu tiles N in chunks of 32 and cannot + # handle a remainder tile smaller than that. + return False, f"requires weight output dim (N={n}) to be a multiple of 32." + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + """Prepack block-quantized FP8 weight for W8A8 AMX BRGEMM. + + For block-quant FP8 models weight is stored as [N, K] (NOT transposed, + unlike non-block FP8 where vLLM stores [K, N]). + + float8_linear_prepack_cpu expects: + weight: [N, K] + scales: [N, G] where G = K / block_k + + weight_scale_inv from the checkpoint has shape [ceil(N/block_n), G], + so we expand rows from ceil(N/block_n) to N via repeat_interleave. + """ + params = self._get_layer_params(layer) + weight = params.weight # [N, K] + + # MLA's kv_b_proj (see `_cpu_skip_gemm_dispatch`): MLA reads this + # weight to derive W_UK/W_UV, so keep the checkpoint [N, K] layout and + # its block scales; apply_weights falls back to the W8A16 op. + if getattr(layer, "_cpu_skip_gemm_dispatch", False): + return + + weight_scale = ( + params.weight_scale_inv + if params.weight_scale_inv is not None + else params.weight_scale + ) # [ceil(N/block_n), ceil(K/block_k)] + assert weight_scale is not None, "FP8 W8A8 block kernel requires weight scale." + + N = weight.shape[0] + block_n = self.weight_group_shape.row # e.g. 128 + + # Expand scale rows: [ceil(N/block_n), G] → [N, G] + if weight_scale.size(0) != N: + weight_scale_expanded = weight_scale.repeat_interleave(block_n, dim=0)[ + :N, : + ].contiguous() + else: + weight_scale_expanded = weight_scale.contiguous() + + packed_weight, packed_scales = torch.ops._C.float8_linear_prepack_cpu( + weight, weight_scale_expanded.to(torch.float32) + ) + + scale_attr = ( + params.WEIGHT_SCALE_INV + if params.weight_scale_inv is not None + else params.WEIGHT_SCALE + ) + replace_parameter( + layer, + params.WEIGHT, + torch.nn.Parameter(packed_weight, requires_grad=False), + ) + replace_parameter( + layer, + scale_attr, + torch.nn.Parameter(packed_scales, requires_grad=False), + ) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + params = self._get_layer_params(layer) + weight_scale = ( + params.weight_scale_inv + if params.weight_scale_inv is not None + else params.weight_scale + ) + + x_2d = x.reshape(-1, x.shape[-1]) if x.dim() > 2 else x + if getattr(layer, "_cpu_skip_gemm_dispatch", False): + # Unpacked weight: fp8_scaled_mm_with_quant has no is_vnni option, + # so fall back to the W8A16 op (BF16 activation, block-scaled FP8 + # weight) for this layer. + out = torch.ops._C.fp8_scaled_mm_cpu( + x_2d, + params.weight, + weight_scale, + list(self.weight_group_shape), + bias, + x.dtype, + False, # is_vnni + ) + return out.reshape(x.shape[:-1] + (out.size(-1),)) if x.dim() > 2 else out + + # act_scales=None + channelwise=True → dynamic per-token activation quant + out = ops.fp8_scaled_mm_with_quant( + x_2d, + None, + True, + params.weight, + weight_scale, + bias, + x.dtype, + ) + return out.reshape(x.shape[:-1] + (out.size(-1),)) if x.dim() > 2 else out + + def apply_block_scaled_mm( + self, + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + ) -> torch.Tensor: + raise NotImplementedError( + "CPUFp8W8A8BlockScaledMMKernel overrides apply_weights directly." + ) + + +class CPUFP8W8A8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): + """FP8 W8A8 GEMM with dynamic per-token activation quantization on CPU.""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.is_cpu(): + return False, "requires CPU platform." + if not torch.cpu._is_amx_tile_supported(): + return False, "requires AMX tile support (Sapphire Rapids or newer)." + if not ops.cpu_has_amx_fp8(): + return False, ( + "requires native AMX-FP8 hardware (Diamond Rapids); falls back " + "to the Torch FP8 kernels otherwise." + ) + return True, None + + @classmethod + def can_implement( + cls, config: FP8ScaledMMLinearLayerConfig + ) -> tuple[bool, str | None]: + if config.out_dtype not in (torch.bfloat16, torch.float16, torch.float32): + return False, "Only bfloat16/float16/float32 output dtype supported." + # Reject block-quantized weights — those go to CPUFp8BlockScaledMMKernel + weight_gs = config.weight_quant_key.scale.group_shape + if weight_gs.col > 0 and weight_gs.row > 0: + return False, "Block-quantized weights not handled by FP8 W8A8 kernel." + n = config.weight_shape[0] + if n % 32 != 0: + # float8_linear_prepack_cpu tiles N in chunks of 32 and cannot + # handle a remainder tile smaller than that. + return False, f"requires weight output dim (N={n}) to be a multiple of 32." + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + """Prepack weights with float8_linear_prepack_cpu (VNNI + block layout).""" + w_name, ws_name, xs_name, _ = self.layer_param_names + weight = getattr(layer, w_name) + weight_scale = getattr(layer, ws_name) + + # vLLM stores linear weight as [K, N], while the CPU prepack op expects [N, K]. + N = weight.shape[-1] + + # Ensure scale is 2D [N, G] + ws = weight_scale + if ws is None: + raise RuntimeError("FP8 W8A8 kernel requires weight scale.") + ws = ws.to(torch.float32) + + # Fused modules may carry per-shard scales (e.g. QKV/MLP shards). + # Expand those to per-output-channel scales before prepack. + if ( + len(getattr(layer, "logical_widths", [])) > 1 + and ws.dim() <= 1 + and ws.numel() == len(layer.logical_widths) + ): + ws = convert_to_channelwise(ws.reshape(-1), layer.logical_widths) + + if ws.dim() == 0: + ws = ws.view(1, 1).expand(N, 1) + elif ws.dim() == 1: + if ws.numel() == 1: + ws = ws.view(1, 1).expand(N, 1) + elif ws.numel() == N: + ws = ws.unsqueeze(1) # [N] -> [N, 1] + else: + raise RuntimeError( + "FP8 W8A8 kernel got unsupported 1D weight scale shape: " + f"{tuple(ws.shape)} for N={N}." + ) + + # Align scale rows to output channels N. + if ws.size(0) != N: + if ws.size(0) == 1: + ws = ws.expand(N, ws.size(1)).contiguous() + else: + raise RuntimeError( + "FP8 W8A8 kernel got mismatched weight scale rows: " + f"ws_rows={ws.size(0)}, expected N={N}." + ) + + w_2d = weight.t().contiguous() # [N, K] + packed_weight, packed_scales = torch.ops._C.float8_linear_prepack_cpu(w_2d, ws) + + replace_parameter( + layer, + w_name, + torch.nn.Parameter(packed_weight, requires_grad=False), + ) + replace_parameter( + layer, + ws_name, + torch.nn.Parameter(packed_scales, requires_grad=False), + ) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + w_name, ws_name, xs_name, _ = self.layer_param_names + weight = getattr(layer, w_name) + weight_scale = getattr(layer, ws_name) + # For static quantization, use pre-computed activation scale + act_scales = getattr(layer, xs_name, None) if xs_name else None + + x_2d = x.reshape(-1, x.shape[-1]) if x.dim() > 2 else x + # channelwise=True: per-token (per-row) quantization when act_scales is None + # channelwise=False: per-tensor quantization when act_scales is provided + channelwise = act_scales is None + out = ops.fp8_scaled_mm_with_quant( + x_2d, + act_scales, + channelwise, + weight, + weight_scale, + bias, + x.dtype, + ) + return out.reshape(x.shape[:-1] + (out.size(-1),)) if x.dim() > 2 else out + + def apply_scaled_mm(self, A, B, out_dtype, As, Bs, bias=None, **kwargs): + raise NotImplementedError( + "CPUFP8W8A8ScaledMMLinearKernel uses apply_weights directly." + ) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 26547c7353..3104ac25c6 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -1055,7 +1055,7 @@ class MLAAttention(nn.Module, AttentionLayerBase): mqa_q_nope, self.impl._w_uk_packed, # type: ignore[attr-defined] True, - None, + self.impl._w_scale, # type: ignore[attr-defined] ) mqa_ql_nope = mqa_ql_nope.transpose(0, 1) else: @@ -1401,7 +1401,7 @@ class MLAAttention(nn.Module, AttentionLayerBase): x, self.impl._w_uv_packed, # type: ignore[attr-defined] True, - None, + self.impl._w_scale, # type: ignore[attr-defined] ) else: # Multiply + Transpose (N, B, L) x (N, L, V)->(N, B, V)->(B, N, V) diff --git a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py index 858b351864..3e7376fe90 100644 --- a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py @@ -15,6 +15,7 @@ from vllm._custom_ops import ( convert_weight_packed_scale_zp, cpu_fused_moe, cpu_fused_moe_int8, + cpu_has_amx_fp8, cpu_prepack_moe_weight, cpu_prepack_moe_weight_int8, fused_experts_cpu, @@ -467,6 +468,7 @@ class CPUExpertsFp8(mk.FusedMoEExpertsModular): RoutingMethodType.Default, RoutingMethodType.Renormalize, RoutingMethodType.RenormalizeNaive, + RoutingMethodType.DeepSeekV3, RoutingMethodType.DeepseekV4, RoutingMethodType.DeepSeekV3, ] @@ -1612,3 +1614,218 @@ class ZenCPUExpertsInt8(mk.FusedMoEExpertsModular): self._w1_scale_bf16, self._w2_scale_bf16, ) + + +# =========================================================================== +# FP8 W8A8 MoE +# =========================================================================== + + +def prepare_fp8_w8a8_moe_layer_for_cpu( + w13: torch.Tensor, + w2: torch.Tensor, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Prepack FP8 W8A8 MoE weights for CPU kernel.""" + E = w13.size(0) + + packed_w13_list = [] + packed_w13_scale_list = [] + + for i in range(E): + w1_exp = w13[i] # [2N, K] FP8 + ws_exp = w13_scale[i] # [2N, K/128] or [2N_blocks, K/128] + + # Align scale rows to weight rows if necessary (for block-128 checkpoints + # where scale may have fewer rows: [2N/128, K/128] → [2N, K/128]) + if ws_exp.size(0) < w1_exp.size(0): + repeat_factor = w1_exp.size(0) // ws_exp.size(0) + ws_exp = torch.repeat_interleave(ws_exp, repeat_factor, dim=0) + ws_exp = ws_exp[: w1_exp.size(0), :].contiguous() + + pw, ps = torch.ops._C.float8_linear_prepack_cpu(w1_exp, ws_exp) + packed_w13_list.append(pw) + packed_w13_scale_list.append(ps) + + packed_w13 = torch.stack(packed_w13_list) + packed_w13_scale = torch.stack(packed_w13_scale_list) + + # w2 uses the W8A16 path (BF16 activation × FP8 weight), VNNI packed + packed_w2 = torch.ops._C.convert_weight_packed(w2) + return packed_w13, packed_w13_scale, packed_w2, w2_scale + + +class CPUExpertsFp8W8A8(mk.FusedMoEExpertsModular): + """CPU FP8 W8A8 block-quantized modular MoE experts.""" + + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + ): + super().__init__(moe_config, quant_config) + + @property + def expects_unquantized_inputs(self) -> bool: + # Receives BF16 hidden states; quantization to FP8 is done inside apply() + return True + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + @staticmethod + def _supports_current_device() -> bool: + # Requires native AMX-FP8 hardware. + return current_platform.is_cpu() and cpu_has_amx_fp8() + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return False + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + return activation == MoEActivation.SILU + + @staticmethod + def _supports_parallel_config( + moe_parallel_config: FusedMoEParallelConfig, + ) -> bool: + return True + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + # Block-quantized FP8. + return (weight_key, activation_key) == ( + kFp8Static128BlockSym, + kFp8Dynamic128Sym, + ) + + @staticmethod + def _supports_routing_method( + routing_method: RoutingMethodType, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return routing_method in [ + RoutingMethodType.Default, + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + RoutingMethodType.DeepSeekV3, + RoutingMethodType.DeepseekV4, + ] + + @staticmethod + def _supports_router_logits_dtype( + router_logits_dtype: torch.dtype | None, + routing_method: RoutingMethodType, + ) -> bool: + return True + + def moe_problem_size( + self, + a1: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_ids: torch.Tensor, + ) -> tuple[int, int, int, int, int]: + # float8_linear_prepack_cpu blocks w1 into [E, Nc, Kc, block_k, block_n], + # not the (E, N, K) layout the base implementation assumes -- N isn't + # recoverable from that shape, so read it from moe_config instead. + E = w1.shape[0] + K = a1.size(-1) + N = ( + self.moe_config.intermediate_size_per_partition + * self.moe_config.w13_num_shards + ) + M = a1.size(0) if a1.dim() == 2 else a1.size(1) + topk = topk_ids.size(1) + return E, M, N, K, topk + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + # fused_experts_cpu manages its own scratch space. + return (0,), (0,), (M, K) + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + return TopKWeightAndReduceNoOP() + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ) -> None: + if apply_router_weight_on_input: + # fused_experts_cpu always applies topk_weights internally on + # combine; MoEPrepareAndFinalizeNoDPEPModular.prepare() would + # also pre-apply it to hidden_states, double-weighting the + # output. Not needed by any CPU FP8 model today. + raise NotImplementedError( + "CPUExpertsFp8W8A8 does not support apply_router_weight_on_input=True." + ) + + block_shape = ( + list(self.quant_config.block_shape) + if self.quant_config.block_shape + else ( + [ + self.quant_config._w1.shape.row, + self.quant_config._w1.shape.col, + ] + if self.quant_config._w1.shape is not None + else None + ) + ) + + # Quantize hidden_states (BF16) → FP8 with per-token scales + hidden_states_2d = hidden_states.reshape(-1, hidden_states.shape[-1]) + x_fp8, x_scales = torch.ops._C.quantize_fp8e4m3_vec( + hidden_states_2d, True, None + ) + + fused_experts_cpu( + output, + x_fp8, + w1, + w2, + topk_weights, + topk_ids, + CPUQuantMethod.FP8_W8A8, + self.w1_scale, # w1_scale + self.w2_scale, # w2_scale + None, # w1_zero + None, # w2_zero + block_shape, # block_size + None, # w1_bias + None, # w2_bias + None, # alpha + None, # limit + True, # is_vnni + x_scales, # a1_scale: per-token FP8 activation scales + ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index 6a6ce34d16..abf0896204 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -57,7 +57,9 @@ class Fp8MoeBackend(Enum): VLLM_CUTLASS = "VLLM_CUTLASS" BATCHED_VLLM_CUTLASS = "BATCHED_VLLM_CUTLASS" XPU = "XPU" + # CPU FP8 W8A16 (BF16 activations); CPU_W8A8 needs native AMX-FP8. CPU = "CPU" + CPU_W8A8 = "CPU_W8A8" HPC = "HPC" # Dequantize-to-BF16 emulation for MXFP8 on devices without a native # MXFP8 MoE kernel (e.g. ROCm). Weights pass through unchanged here. @@ -92,6 +94,7 @@ def _get_priority_backends( Fp8MoeBackend.BATCHED_VLLM_CUTLASS, Fp8MoeBackend.BATCHED_TRITON, Fp8MoeBackend.XPU, + Fp8MoeBackend.CPU_W8A8, Fp8MoeBackend.CPU, Fp8MoeBackend.HPC, ] @@ -129,8 +132,10 @@ def _get_priority_backends( _move_to_front(_AVAILABLE_BACKENDS, Fp8MoeBackend.XPU) if current_platform.is_cpu(): - # CPU platform uses FP8 W8A16 fused MoE kernel. + # W8A8 first: it falls through to the W8A16 backend whenever the + # hardware (AMX-FP8) or the config isn't supported. _move_to_front(_AVAILABLE_BACKENDS, Fp8MoeBackend.CPU) + _move_to_front(_AVAILABLE_BACKENDS, Fp8MoeBackend.CPU_W8A8) return _AVAILABLE_BACKENDS @@ -238,6 +243,13 @@ def backend_to_kernel_cls( return [CPUExpertsFp8] + elif backend == Fp8MoeBackend.CPU_W8A8: + from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( + CPUExpertsFp8W8A8, + ) + + return [CPUExpertsFp8W8A8] + elif backend == Fp8MoeBackend.HPC: from vllm.model_executor.layers.fused_moe.hpc_moe import ( HPCExperts, @@ -652,6 +664,14 @@ def convert_to_fp8_moe_kernel_format( ) w13, w2 = prepare_fp8_moe_layer_for_cpu(w13, w2) + elif fp8_backend == Fp8MoeBackend.CPU_W8A8: + from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( + prepare_fp8_w8a8_moe_layer_for_cpu, + ) + + w13, w13_scale, w2, w2_scale = prepare_fp8_w8a8_moe_layer_for_cpu( + w13, w2, w13_scale, w2_scale + ) else: if fp8_backend not in [ Fp8MoeBackend.TRITON, @@ -697,7 +717,15 @@ def make_fp8_moe_quant_config( In a future PR, we will have this function should be a method of the modular kernel itself. """ - # MARLIN and CPU are mixed precision W8A16 config. + if fp8_backend == Fp8MoeBackend.CPU_W8A8: + return fp8_w8a8_moe_quant_config( + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + block_shape=block_shape, + ) + + # MARLIN and CPU (W8A16) are mixed precision W8A16 configs. if fp8_backend == Fp8MoeBackend.MARLIN or fp8_backend == Fp8MoeBackend.CPU: return fp8_w8a16_moe_quant_config( w1_scale=w1_scale, diff --git a/vllm/model_executor/layers/quantization/online/fp8.py b/vllm/model_executor/layers/quantization/online/fp8.py index 2a8ef0b69f..ed018c81bf 100644 --- a/vllm/model_executor/layers/quantization/online/fp8.py +++ b/vllm/model_executor/layers/quantization/online/fp8.py @@ -737,11 +737,13 @@ class Fp8PtpcOnlineMoEMethod(_Fp8OnlineMoEBase): ) # Reject backends whose make_fp8_moe_quant_config branch silently # drops per_act_token_quant / per_out_ch_quant or collapses scales: - # MARLIN / CPU route through fp8_w8a16_moe_quant_config; FLASHINFER_* - # fold scales into a per-tensor alpha (oracle/fp8.py). + # MARLIN / CPU route through fp8_w8a16_moe_quant_config, CPU_W8A8 + # ignores both flags, and FLASHINFER_* fold scales into a per-tensor + # alpha (oracle/fp8.py). if self.fp8_backend in ( Fp8MoeBackend.MARLIN, Fp8MoeBackend.CPU, + Fp8MoeBackend.CPU_W8A8, Fp8MoeBackend.FLASHINFER_CUTLASS, Fp8MoeBackend.FLASHINFER_TRTLLM, ): diff --git a/vllm/v1/attention/backends/mla/amx_mla.py b/vllm/v1/attention/backends/mla/amx_mla.py index 5f9ac7f17a..38c8b68af9 100644 --- a/vllm/v1/attention/backends/mla/amx_mla.py +++ b/vllm/v1/attention/backends/mla/amx_mla.py @@ -61,6 +61,44 @@ _MIN_WORK_PER_SPLIT = 512 _SPLIT_OCCUPANCY_MULTIPLIER = 2 +def _block_quant_to_tensor_quant( + weight: torch.Tensor, # FP8 [N, K] block-quantized weight + weight_scale: torch.Tensor, # [ceil(N/bn), ceil(K/bk)] block scales + block_size: list[int], # [block_n, block_k] +) -> tuple[torch.Tensor, torch.Tensor]: + """Re-quantize FP8 block-quant weight to FP8 per-tensor quant. + + Returns (fp8_weight, dequant_scale) where fp8_weight * dequant_scale + approximates the original float weight. dequant_scale is a 0-dim + scalar tensor as expected by bmm_cpu's ``scale`` argument. + + NOTE: block->tensor re-quantization may slightly affect accuracy + (equivalent to SGLang's block_quant_to_tensor_quant for CPU). + """ + block_n, block_k = block_size[0], block_size[1] + n, k = weight.shape + + # Expand block scales [ceil(N/bn), ceil(K/bk)] -> [N, K] + scale = weight_scale.to(torch.float32) + if scale.shape[0] < n: + scale = scale.repeat_interleave(block_n, dim=0)[:n, :] + if scale.shape[1] < k: + scale = scale.repeat_interleave(block_k, dim=1)[:, :k] + + # Dequantize to float32 + x_f32 = weight.to(torch.float32) * scale + + # Re-quantize per-tensor back to FP8 + fp8_max = torch.finfo(weight.dtype).max # 448.0 for e4m3fn + amax = x_f32.abs().amax().clamp(min=1e-12) + quant_scale = fp8_max / amax + x_fp8 = (x_f32 * quant_scale).clamp(-fp8_max, fp8_max).to(weight.dtype) + + # dequant_scale = 1 / quant_scale (0-dim tensor for bmm_cpu) + dequant_scale = (amax / fp8_max).to(torch.float32) + return x_fp8, dequant_scale + + def _compute_num_kv_splits(max_seq_len: int, num_threads: int) -> int: """Mirrors TritonMLAImpl's _compute_num_kv_splits, using the CPU thread count in place of SM count.""" @@ -208,6 +246,9 @@ class AMXMLAImpl(MLACommonImpl[MLACommonMetadata]): ) self._w_uk_packed: torch.Tensor | None = None self._w_uv_packed: torch.Tensor | None = None + self._w_scale: torch.Tensor | None = ( + None # per-tensor dequant scale (FP8 W8A16) + ) def process_weights_after_loading(self, act_dtype: torch.dtype) -> None: # forward_mha (unlike forward_mqa) receives no `layer` argument, so it @@ -216,12 +257,40 @@ class AMXMLAImpl(MLACommonImpl[MLACommonMetadata]): # impl derives its own copies from self.kv_b_proj (already available # via the constructor) and VNNI-packs them once, ahead of time, for # use with bmm_cpu in forward_mha. - kv_b_proj_weight = get_and_maybe_dequant_weights( - self.kv_b_proj, out_dtype=act_dtype - ).T - kv_b_proj_weight = kv_b_proj_weight.view( - self.kv_lora_rank, self.num_heads, self.qk_nope_head_dim + self.v_head_dim - ) + # + # kv_b_proj carries `_cpu_skip_gemm_dispatch`, so the CPU FP8 kernels + # leave its weight in the checkpoint [N, K] layout. For block-quant + # FP8 checkpoints, re-quantize per-tensor and keep the BMM in FP8 + # (W8A16); everything else goes through the BF16 path. + weight = getattr(self.kv_b_proj, "weight", None) + block_size = getattr(self.kv_b_proj, "weight_block_size", None) + if ( + weight is not None + and weight.dtype == torch.float8_e4m3fn + and block_size is not None + ): + block_scale = getattr(self.kv_b_proj, "weight_scale_inv", None) + if block_scale is None: + block_scale = self.kv_b_proj.weight_scale + fp8_weight, self._w_scale = _block_quant_to_tensor_quant( + weight, block_scale, block_size + ) + # [N, K] -> [K, N] = [kv_lora_rank, heads * (nope + v)] + kv_b_proj_weight = fp8_weight.T.contiguous().view( + self.kv_lora_rank, + self.num_heads, + self.qk_nope_head_dim + self.v_head_dim, + ) + else: + kv_b_proj_weight = get_and_maybe_dequant_weights( + self.kv_b_proj, out_dtype=act_dtype + ).T + kv_b_proj_weight = kv_b_proj_weight.view( + self.kv_lora_rank, + self.num_heads, + self.qk_nope_head_dim + self.v_head_dim, + ) + self._w_scale = None w_uk, w_uv = kv_b_proj_weight.split( [self.qk_nope_head_dim, self.v_head_dim], dim=-1 ) @@ -348,7 +417,9 @@ class AMXMLAImpl(MLACommonImpl[MLACommonMetadata]): ql_nope = torch.empty( num_tokens, num_heads, self.kv_lora_rank, dtype=q.dtype, device=q.device ) - ops.bmm_cpu(ql_nope.transpose(0, 1), q_nope_t, self._w_uk_packed, True, None) + ops.bmm_cpu( + ql_nope.transpose(0, 1), q_nope_t, self._w_uk_packed, True, self._w_scale + ) mqa_q = torch.cat([ql_nope, q_pe], dim=-1) kv_cache_flat = kv_c_and_k_pe_cache.view(-1, 1, self.head_size) @@ -402,4 +473,4 @@ class AMXMLAImpl(MLACommonImpl[MLACommonMetadata]): output_view = output.view(num_tokens, num_heads, self.v_head_dim).transpose( 0, 1 ) - ops.bmm_cpu(output_view, attn_out_t, self._w_uv_packed, True, None) + ops.bmm_cpu(output_view, attn_out_t, self._w_uv_packed, True, self._w_scale)