[CPU] Add CPU FP8 W8A8 linear/MoE support (#49942)
Signed-off-by: yuwenzho <yuwen.zhou@intel.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>
This commit is contained in:
Yuwen Zhou
co-authored by
Li, Jiang <jiang1.li@intel.com>
Li, Jiang <jiang1.li@intel.com>
parent
6df2b1a8c6
commit
b346479065
@@ -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"
|
||||
|
||||
+2
-11
@@ -1,16 +1,7 @@
|
||||
#include "cpu_attn_dispatch_generated.h"
|
||||
|
||||
#if defined(__x86_64__) || defined(_M_X64)
|
||||
#include <cpuid.h>
|
||||
// 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.
|
||||
|
||||
@@ -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 <cpuid.h>
|
||||
|
||||
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
|
||||
@@ -12,11 +12,45 @@
|
||||
#include <ATen/ATen.h>
|
||||
#include <ATen/Dispatch.h>
|
||||
#include <ATen/Parallel.h>
|
||||
#include <cstring>
|
||||
|
||||
#if defined(_OPENMP)
|
||||
#include <omp.h>
|
||||
#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<at::Tensor>& opt) {
|
||||
return opt.has_value() ? opt.value().data_ptr<T>() : 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<bool, int64_t, int64_t, int64_t> 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 <typename T>
|
||||
void zero_buffer(T* data, int64_t size) {
|
||||
const int32_t vec_size = at::vec::Vectorized<T>::size();
|
||||
auto zero_vec = at::vec::Vectorized<T>(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 <typename T>
|
||||
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
|
||||
|
||||
@@ -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<int64_t>(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<at::Tensor, at::Tensor> quantize_fp8e4m3_vec(
|
||||
const at::Tensor& t,
|
||||
bool channelwise,
|
||||
c10::optional<at::Tensor> scale_opt);
|
||||
|
||||
// FP8 W8A8 fused MoE kernel (defined in moe_fp8_w8a8.cpp)
|
||||
template <typename scalar_t>
|
||||
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);
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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<CPUQuantMethod::FP8_W8A16>(w1_scale, w2_scale, block_size);
|
||||
} else if (moe_comp_method == CPUQuantMethod::MXFP4) {
|
||||
check_moe_scales<CPUQuantMethod::MXFP4>(w1_scale, w2_scale, block_size);
|
||||
} else if (moe_comp_method == CPUQuantMethod::FP8_W8A8) {
|
||||
check_moe_scales<CPUQuantMethod::FP8_W8A8>(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<at::Tensor>& w2_scale,
|
||||
const std::optional<at::Tensor>& w1_zero,
|
||||
const std::optional<at::Tensor>& w2_zero,
|
||||
const std::optional<at::Tensor>& a1_scale,
|
||||
const std::optional<std::vector<int64_t>> block_size,
|
||||
const std::optional<at::Tensor>& w1_bias,
|
||||
const std::optional<at::Tensor>& 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<scalar_t>(),
|
||||
intermediate_cache0,
|
||||
intermediate_cache1,
|
||||
intermediate_cache2,
|
||||
A_tmp,
|
||||
B_tmp,
|
||||
C_tmp,
|
||||
Ukernel_tmp,
|
||||
hidden_states.data_ptr<at::Float8_e4m3fn>(),
|
||||
packed_w1.data_ptr<at::Float8_e4m3fn>(),
|
||||
packed_w2.data_ptr<at::Float8_e4m3fn>(),
|
||||
a1_scale.has_value() ? a1_scale.value().data_ptr<float>() : nullptr,
|
||||
w1s.data_ptr<float>(),
|
||||
w2s.data_ptr<float>(),
|
||||
block_size_N,
|
||||
block_size_K,
|
||||
topk_weights_.data_ptr<float>(),
|
||||
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));
|
||||
|
||||
@@ -0,0 +1,432 @@
|
||||
// Adapted from
|
||||
// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
|
||||
|
||||
// clang-format off
|
||||
|
||||
#include <c10/util/Unroll.h>
|
||||
|
||||
#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 <typename out_dtype, int64_t N, int act_quant_mode, int wei_quant_mode>
|
||||
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<N_UNROLL>{}([&](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<out_dtype, float>::value) {
|
||||
_mm512_storeu_ps(c_ptr + i * lda + j, y_vec);
|
||||
} else if constexpr (std::is_same<out_dtype, at::BFloat16>::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<out_dtype, at::Half>::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<out_dtype>(y_buf[i * N + j] * a_scale * b_scale);
|
||||
}
|
||||
} // for M
|
||||
}
|
||||
|
||||
#else // no AVX512
|
||||
|
||||
template <typename out_dtype, int64_t N, int act_quant_mode, int wei_quant_mode>
|
||||
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<out_dtype>(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<scalar_t>)
|
||||
// Stage 3: out = weighted_sum(ic2, topk_weights) [M, K]
|
||||
|
||||
template <typename scalar_t>
|
||||
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<at::BFloat16*>(
|
||||
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<scalar_t, BLOCK_N_VAL, PER_ROW, PER_GROUP>(
|
||||
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<scalar_t, BLOCK_N_VAL, PER_ROW, PER_GROUP>(
|
||||
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<at::Float8_e4m3fn>(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<at::Float8_e4m3fn>(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<scalar_t>(
|
||||
/* 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>( \
|
||||
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);
|
||||
@@ -73,6 +73,7 @@ void fused_experts_cpu(at::Tensor& out, at::Tensor& hidden_states,
|
||||
const std::optional<at::Tensor>& w2_scale,
|
||||
const std::optional<at::Tensor>& w1_zero,
|
||||
const std::optional<at::Tensor>& w2_zero,
|
||||
const std::optional<at::Tensor>& a1_scale,
|
||||
const std::optional<std::vector<int64_t>> block_size,
|
||||
const std::optional<at::Tensor>& w1_bias,
|
||||
const std::optional<at::Tensor>& w2_bias,
|
||||
@@ -91,6 +92,18 @@ at::Tensor fp8_scaled_mm_cpu(at::Tensor& mat1, at::Tensor& mat2,
|
||||
const std::optional<at::Tensor>& bias,
|
||||
at::ScalarType out_dtype, bool is_vnni);
|
||||
|
||||
// Adapted from sglang: FP8 W8A8 kernels
|
||||
std::tuple<at::Tensor, at::Tensor> 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<at::Tensor>& act_scales,
|
||||
bool channelwise, const at::Tensor& weight,
|
||||
const at::Tensor& weight_scales,
|
||||
const std::optional<at::Tensor>& bias,
|
||||
at::ScalarType output_dtype);
|
||||
std::tuple<at::Tensor, at::Tensor> quantize_fp8e4m3_vec(
|
||||
const at::Tensor& t, bool channelwise, c10::optional<at::Tensor> 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(
|
||||
|
||||
@@ -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"]))
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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: [
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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."
|
||||
)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
):
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user