@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 88a860d754
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 3, 10:30 UTC

[CPU] Add MXFP4 W4A16 MoE support (#41922)

Signed-off-by: yuwenzho <yuwen.zhou@intel.com> Signed-off-by: Yuwen Zhou <yuwen.zhou@intel.com>

Yuwen Zhou committed May 18, 2026 at 18:04 UTC 88a860d7545aad69661daad7a1c2b04f59c76144
16 files changed +1173 -350
.buildkite/hardware_tests/cpu.yaml
+3 -3
@@ -12,7 +12,7 @@ steps:
12 - vllm/_custom_ops.py
13 - tests/kernels/attention/test_cpu_attn.py
14 - tests/kernels/moe/test_cpu_fused_moe.py
15 - - tests/kernels/moe/test_cpu_fp8_fused_moe.py
15 + - tests/kernels/moe/test_cpu_quant_fused_moe.py
16 - tests/kernels/test_onednn.py
17 - tests/kernels/test_awq_int4_to_int8.py
18 - tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
@@ -21,7 +21,7 @@ steps:
21 bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
22 pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
23 pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
24 - pytest -x -v -s tests/kernels/moe/test_cpu_fp8_fused_moe.py
24 + pytest -x -v -s tests/kernels/moe/test_cpu_quant_fused_moe.py
25 pytest -x -v -s tests/kernels/test_onednn.py
26 pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py
27 pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py"
@@ -70,7 +70,7 @@ steps:
70 - tests/quantization/test_cpu_wna16.py
71 commands:
72 - |
73 - bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 20m "
73 + bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
74 pytest -x -v -s tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs
75 pytest -x -v -s tests/quantization/test_cpu_wna16.py"
76
csrc/cpu/sgl-kernels/common.h
+7
@@ -348,6 +348,13 @@ inline int get_cache_blocks<at::Float8_e4m3fn>(int chunk_size) {
348 return std::min(MAX_CACHE_BLOCK_SIZE, cache_block_size);
349 }
350
351 +template <>
352 +inline int get_cache_blocks<uint8_t>(int chunk_size) {
353 + // mxfp4 uses bf16 as accumulate type
354 + int cache_block_size = get_cache_blocks<at::BFloat16>(chunk_size);
355 + return std::min(MAX_CACHE_BLOCK_SIZE, cache_block_size);
356 +}
357 +
358 // 2d sequential loop in range : [mb0, mb1), [nb0, nb1)
359 template <typename T, typename func_t>
360 inline void loop_2d(int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1, int64_t chunk_size, const func_t& f) {
csrc/cpu/sgl-kernels/gemm.h
+49 -9
@@ -72,7 +72,17 @@ inline int64_t get_row_size(int64_t K, bool use_int8_w8a8) {
72 return use_int8_w8a8 ? K + sizeof(int32_t) : K;
73 }
74
75 -enum class CPUQuantMethod : int64_t { BF16 = 0, INT8_W8A8 = 1, FP8_W8A16 = 2, INT4_W4A8 = 3 };
75 +enum class CPUAcTMethod : int { silu_and_mul = 0, swiglu = 1 };
76 +
77 +constexpr bool operator==(CPUAcTMethod a, int b) {
78 + return static_cast<int>(a) == b;
79 +}
80 +
81 +constexpr bool operator==(int a, CPUAcTMethod b) {
82 + return a == static_cast<int>(b);
83 +}
84 +
85 +enum class CPUQuantMethod : int64_t { BF16 = 0, INT8_W8A8 = 1, FP8_W8A16 = 2, INT4_W4A8 = 3, MXFP4 = 4 };
86
87 constexpr bool operator==(CPUQuantMethod a, int64_t b) {
88 return static_cast<int64_t>(a) == b;
@@ -99,6 +109,9 @@ inline int64_t get_4bit_block_k_size(int64_t group_size) {
109 // pack weight to vnni format
110 at::Tensor convert_weight_packed(at::Tensor& weight);
111
112 +// pack scale to blocked format for mxfp4
113 +at::Tensor convert_scale_packed(at::Tensor& scale);
114 +
115 // pack weight to vnni format for int4
116 std::tuple<at::Tensor, at::Tensor, at::Tensor>
117 convert_weight_packed_scale_zp(at::Tensor qweight, at::Tensor qzeros, at::Tensor scales);
@@ -129,9 +142,9 @@ void fused_experts_int8_kernel_impl(
142 int64_t topk,
143 int64_t num_tokens_post_pad);
144
132 -// moe implementations for fp8 w8a16
133 -template <typename scalar_t>
134 -void fused_experts_fp8_kernel_impl(
145 +// moe implementations for fp8 w8a16 and mxfp4
146 +template <typename scalar_t, typename packed_t, typename param_t, bool is_mxfp4>
147 +void fused_experts_fp_kernel_impl(
148 scalar_t* __restrict__ output,
149 scalar_t* __restrict__ ic0,
150 scalar_t* __restrict__ ic1,
@@ -140,10 +153,12 @@ void fused_experts_fp8_kernel_impl(
153 scalar_t* __restrict__ B_tmp,
154 float* __restrict__ C_tmp,
155 const scalar_t* __restrict__ input,
143 - const at::Float8_e4m3fn* __restrict__ packed_w1,
144 - const at::Float8_e4m3fn* __restrict__ packed_w2,
145 - const float* __restrict__ w1s,
146 - const float* __restrict__ w2s,
156 + const packed_t* __restrict__ packed_w1,
157 + const packed_t* __restrict__ packed_w2,
158 + const float* __restrict__ w1_bias,
159 + const float* __restrict__ w2_bias,
160 + const param_t* __restrict__ w1s,
161 + const param_t* __restrict__ w2s,
162 int64_t block_size_N,
163 int64_t block_size_K,
164 const float* __restrict__ topk_weights,
@@ -155,7 +170,11 @@ void fused_experts_fp8_kernel_impl(
170 int64_t K,
171 int64_t E,
172 int64_t topk,
158 - int64_t num_tokens_post_pad);
173 + int64_t num_tokens_post_pad,
174 + float alpha,
175 + float limit,
176 + CPUAcTMethod act_func,
177 + bool with_bias);
178
179 // shared expert implementation for int8 w8a8
180 template <typename scalar_t>
@@ -266,6 +285,7 @@ void tinygemm_kernel(
285 scalar_t* __restrict__ C,
286 scalar_t* __restrict__ Btmp,
287 float* __restrict__ Ctmp,
288 + const float* __restrict__ Bbias,
289 const float* __restrict__ scale,
290 int64_t M,
291 int64_t N,
@@ -294,6 +314,26 @@ void tinygemm_kernel(
314 int64_t ldc,
315 bool brg);
316
317 +// mxfp4
318 +template <typename scalar_t>
319 +void tinygemm_kernel(
320 + const scalar_t* __restrict__ A,
321 + const uint8_t* __restrict__ B,
322 + scalar_t* __restrict__ C,
323 + scalar_t* __restrict__ Btmp,
324 + float* __restrict__ Ctmp,
325 + const float* __restrict__ Bbias,
326 + const uint8_t* __restrict__ scale,
327 + int64_t M,
328 + int64_t N,
329 + int64_t K,
330 + int64_t lda,
331 + int64_t ldb,
332 + int64_t ldc,
333 + bool brg,
334 + int64_t block_size_K,
335 + bool do_unpack = true);
336 +
337 template <typename scalar_t>
338 void tinygemm_kernel(
339 scalar_t* C,
csrc/cpu/sgl-kernels/gemm_fp8.cpp
+83 -3
@@ -67,6 +67,23 @@ inline void copy_mul_stub(scalar_t* __restrict__ out, const float* __restrict__
67 }
68 }
69
70 +template <>
71 +inline void
72 +copy_add_stub(float* __restrict__ out, const float* __restrict__ input, const float* __restrict__ bias, int64_t size) {
73 + using fVec = at::vec::Vectorized<float>;
74 + constexpr int kVecSize = fVec::size();
75 +
76 + int64_t d;
77 +#pragma GCC unroll 4
78 + for (d = 0; d <= size - kVecSize; d += kVecSize) {
79 + fVec data = fVec::loadu(input + d) + fVec::loadu(bias + d);
80 + data.store(out + d);
81 + }
82 + for (; d < size; ++d) {
83 + out[d] = input[d] + bias[d];
84 + }
85 +}
86 +
87 inline void unpack_B(
88 at::BFloat16* __restrict__ Btmp,
89 const at::Float8_e4m3fn* __restrict__ packed_B,
@@ -343,7 +360,6 @@ struct tinygemm_kernel_nn<at::BFloat16, at::Float8_e4m3fn, float, has_bias, BLOC
360 Unroll<ROWS * COLS>{}(storec);
361 }
362 };
346 -
363 template <int BLOCK_M, int BLOCK_N>
364 struct tinygemm_kernel_nn2<at::BFloat16, BLOCK_M, BLOCK_N> {
365 static inline void apply(
@@ -918,6 +934,7 @@ void tinygemm_kernel(
934 scalar_t* __restrict__ C,
935 scalar_t* __restrict__ Btmp,
936 float* __restrict__ Ctmp,
937 + const float* __restrict__ Bbias,
938 const float* __restrict__ scale,
939 int64_t M,
940 int64_t N,
@@ -928,6 +945,11 @@ void tinygemm_kernel(
945 bool brg,
946 int64_t block_size_K,
947 bool do_unpack) {
948 + if (Bbias != nullptr) {
949 + tinygemm_kernel<scalar_t, at::Float8_e4m3fn, float, true>(
950 + A, B, C, Btmp, Ctmp, scale, Bbias, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
951 + return;
952 + }
953 tinygemm_kernel<scalar_t, at::Float8_e4m3fn, float, false>(
954 A, B, C, Btmp, Ctmp, scale, nullptr, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
955 }
@@ -949,7 +971,6 @@ void tinygemm_kernel(
971 bool brg) {
972 tinygemm_kernel2<scalar_t>(A, B, C, Btmp, Ctmp, scale, M, N, K, lda, ldb, ldc, brg);
973 }
952 -
974 template <typename scalar_t>
975 void tinygemm_kernel(
976 const scalar_t* __restrict__ A,
@@ -957,6 +978,7 @@ void tinygemm_kernel(
978 scalar_t* __restrict__ C,
979 scalar_t* __restrict__ Btmp,
980 float* __restrict__ Ctmp,
981 + const float* __restrict__ Bbias,
982 const uint8_t* __restrict__ scale,
983 int64_t M,
984 int64_t N,
@@ -967,10 +989,68 @@ void tinygemm_kernel(
989 bool brg,
990 int64_t block_size_K,
991 bool do_unpack) {
992 + if (Bbias != nullptr) {
993 + tinygemm_kernel<scalar_t, uint8_t, uint8_t, true>(
994 + A, B, C, Btmp, Ctmp, scale, Bbias, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
995 + return;
996 + }
997 tinygemm_kernel<scalar_t, uint8_t, uint8_t, false>(
998 A, B, C, Btmp, Ctmp, scale, nullptr, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
999 }
1000
1001 +// tinygemm interface
1002 +template <typename scalar_t>
1003 +void tinygemm_kernel(
1004 + const scalar_t* __restrict__ A,
1005 + const at::Float8_e4m3fn* __restrict__ B,
1006 + float* __restrict__ C,
1007 + scalar_t* __restrict__ Btmp,
1008 + const float* __restrict__ Bbias,
1009 + const float* __restrict__ scale,
1010 + int64_t M,
1011 + int64_t N,
1012 + int64_t K,
1013 + int64_t lda,
1014 + int64_t ldb,
1015 + int64_t ldc,
1016 + bool brg,
1017 + int64_t block_size_K,
1018 + bool do_unpack) {
1019 + if (Bbias != nullptr) {
1020 + tinygemm_kernel<scalar_t, at::Float8_e4m3fn, float, true>(
1021 + A, B, C, Btmp, scale, Bbias, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
1022 + return;
1023 + }
1024 + tinygemm_kernel<scalar_t, at::Float8_e4m3fn, float, false>(
1025 + A, B, C, Btmp, scale, nullptr, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
1026 +}
1027 +
1028 +template <typename scalar_t>
1029 +void tinygemm_kernel(
1030 + const scalar_t* __restrict__ A,
1031 + const uint8_t* __restrict__ B,
1032 + float* __restrict__ C,
1033 + scalar_t* __restrict__ Btmp,
1034 + const float* __restrict__ Bbias,
1035 + const uint8_t* __restrict__ scale,
1036 + int64_t M,
1037 + int64_t N,
1038 + int64_t K,
1039 + int64_t lda,
1040 + int64_t ldb,
1041 + int64_t ldc,
1042 + bool brg,
1043 + int64_t block_size_K,
1044 + bool do_unpack) {
1045 + if (Bbias != nullptr) {
1046 + tinygemm_kernel<scalar_t, uint8_t, uint8_t, true>(
1047 + A, B, C, Btmp, scale, Bbias, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
1048 + return;
1049 + }
1050 + tinygemm_kernel<scalar_t, uint8_t, uint8_t, false>(
1051 + A, B, C, Btmp, scale, nullptr, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
1052 +}
1053 +
1054 #define INSTANTIATE_TINYGEMM_TEMPLATE(TYPE_A, TYPE_B, TYPE_S) \
1055 template void tinygemm_kernel<TYPE_A>( \
1056 const TYPE_A* __restrict__ A, \
@@ -978,6 +1058,7 @@ void tinygemm_kernel(
1058 TYPE_A* __restrict__ C, \
1059 TYPE_A* __restrict__ Btmp, \
1060 float* __restrict__ Ctmp, \
1061 + const float* __restrict__ Bbias, \
1062 const TYPE_S* __restrict__ scale, \
1063 int64_t M, \
1064 int64_t N, \
@@ -1020,7 +1101,6 @@ inline const float* get_bias_data(const std::optional<at::Tensor>& bias, int64_t
1101 }
1102 return nullptr;
1103 }
1023 -
1104 // FP8 and MXFP4 WoQ uses the same pattern:
1105 // Btmp : [T, BLOCK_N * K]
1106 // Ctmp : [T, BLOCK_M * BLOCK_N]
csrc/cpu/sgl-kernels/moe.cpp
+78 -9
@@ -812,6 +812,7 @@ void shared_expert_kernel_impl(
812 static inline void check_moe_scales(
813 bool use_int8_w8a8,
814 bool use_fp8_w8a16,
815 + bool use_mxfp4,
816 const std::optional<at::Tensor>& w1_scale,
817 const std::optional<at::Tensor>& w2_scale,
818 const std::optional<std::vector<int64_t>> block_size) {
@@ -825,6 +826,12 @@ static inline void check_moe_scales(
826 TORCH_CHECK(block_size.has_value(), "missing block_size for fp8 w8a16.");
827 TORCH_CHECK(block_size.value().size() == 2, "expect block_size.size() to be 2.");
828 }
829 + if (use_mxfp4) {
830 + TORCH_CHECK(w1_scale.has_value(), "missing w1_scale for mxfp4.");
831 + TORCH_CHECK(w2_scale.has_value(), "missing w2_scale for mxfp4.");
832 + TORCH_CHECK(w1_scale.value().scalar_type() == at::kByte, "expect w1_scale to be uint8.");
833 + TORCH_CHECK(w2_scale.value().scalar_type() == at::kByte, "expect w2_scale to be uint8.");
834 + }
835 }
836
837 #define CHECK_MOE_SCALES_FP8(DIM0, DIM1) \
@@ -839,8 +846,8 @@ static inline void check_moe_scales(
846 TORCH_CHECK(w2s.size(DIM1) == div_up(N, block_size_K))
847
848 // hidden_states: [M, K]
842 -// w1: [E, 2N, K]
843 -// w2: [E, K, N]
849 +// w1: [E, 2N, K] or [E, 2N, K / 2] for uint8
850 +// w2: [E, K, N] or [E, K, N / 2] for uint8
851 // topk_weights: [M, topk]
852 // topk_ids: [M, topk] (int32_t)
853 //
@@ -858,6 +865,10 @@ at::Tensor fused_experts_cpu(
865 const std::optional<at::Tensor>& w1_zero,
866 const std::optional<at::Tensor>& w2_zero,
867 const std::optional<std::vector<int64_t>> block_size,
868 + const std::optional<at::Tensor>& w1_bias,
869 + const std::optional<at::Tensor>& w2_bias,
870 + const std::optional<double>& alpha,
871 + const std::optional<double>& limit,
872 bool is_vnni) {
873 auto packed_w1 = is_vnni ? w1 : convert_weight_packed(w1);
874 auto packed_w2 = is_vnni ? w2 : convert_weight_packed(w2);
@@ -897,8 +908,12 @@ at::Tensor fused_experts_cpu(
908 int64_t topk = topk_weights_.size(1);
909
910 // we use int32_t compensation for int8 w8a8
900 - int64_t packed_K = get_row_size(K, moe_comp_method == CPUQuantMethod::INT8_W8A8);
901 - int64_t packed_N = get_row_size(N, moe_comp_method == CPUQuantMethod::INT8_W8A8);
911 + int64_t packed_K = moe_comp_method == CPUQuantMethod::MXFP4
912 + ? get_row_size<uint8_t>(K)
913 + : get_row_size(K, moe_comp_method == CPUQuantMethod::INT8_W8A8);
914 + int64_t packed_N = moe_comp_method == CPUQuantMethod::MXFP4
915 + ? get_row_size<uint8_t>(N)
916 + : get_row_size(N, moe_comp_method == CPUQuantMethod::INT8_W8A8);
917
918 // check weight shapes
919 CHECK_EQ(w2.size(0), E);
@@ -911,6 +926,7 @@ at::Tensor fused_experts_cpu(
926 check_moe_scales(
927 moe_comp_method == CPUQuantMethod::INT8_W8A8,
928 moe_comp_method == CPUQuantMethod::FP8_W8A16,
929 + moe_comp_method == CPUQuantMethod::MXFP4,
930 w1_scale,
931 w2_scale,
932 block_size);
@@ -965,7 +981,7 @@ at::Tensor fused_experts_cpu(
981 // 5. Aq_tmp : [M, K] or [M * topk, N]
982 // 6. As_tmp : [M * topk]
983 //
968 - // for fp8 w8a16:
984 + // for fp8 w8a16 and mxfp4:
985 // 7. intermediate_cache0 : [M * topk, 2N]
986 // 8. B_tmp : [T, MAX_CACHE_BLOCK_SIZE, BLOCK_N, std::max(K, N)]
987 //
@@ -978,7 +994,7 @@ at::Tensor fused_experts_cpu(
994 if (moe_comp_method == CPUQuantMethod::INT8_W8A8) {
995 buffer_size_nbytes += std::max(M * K, M * topk * N) + M * topk * sizeof(float);
996 }
981 - if (moe_comp_method == CPUQuantMethod::FP8_W8A16) {
997 + if (moe_comp_method == CPUQuantMethod::FP8_W8A16 || moe_comp_method == CPUQuantMethod::MXFP4) {
998 buffer_size_nbytes += M * topk * 2 * N * 2 + num_threads * MAX_CACHE_BLOCK_SIZE * BLOCK_N * std::max(K, N) * 2;
999 }
1000 if (moe_comp_method == CPUQuantMethod::INT4_W4A8) {
@@ -1031,9 +1047,11 @@ at::Tensor fused_experts_cpu(
1047 float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K));
1048 scalar_t* __restrict__ intermediate_cache0 = (scalar_t*)((void*)(C_tmp + num_threads * 2 * BLOCK_M * BLOCK_N));
1049 scalar_t* __restrict__ B_tmp = (scalar_t*)((void*)(intermediate_cache0 + M * topk * 2 * N));
1050 + bool with_bias = w1_bias.has_value();
1051 + auto act_func = alpha.has_value() && limit.has_value() ? CPUAcTMethod::swiglu : CPUAcTMethod::silu_and_mul;
1052
1053 CHECK_MOE_SCALES_FP8(1, 2);
1036 - fused_experts_fp8_kernel_impl(
1054 + fused_experts_fp_kernel_impl<scalar_t, at::Float8_e4m3fn, float, false>(
1055 out_hidden_states.data_ptr<scalar_t>(),
1056 intermediate_cache0,
1057 intermediate_cache1,
@@ -1044,6 +1062,8 @@ at::Tensor fused_experts_cpu(
1062 hidden_states.data_ptr<scalar_t>(),
1063 packed_w1.data_ptr<at::Float8_e4m3fn>(),
1064 packed_w2.data_ptr<at::Float8_e4m3fn>(),
1065 + with_bias ? w1_bias.value().data_ptr<float>() : nullptr,
1066 + with_bias ? w2_bias.value().data_ptr<float>() : nullptr,
1067 w1s.data_ptr<float>(),
1068 w2s.data_ptr<float>(),
1069 block_size_N,
@@ -1057,7 +1077,56 @@ at::Tensor fused_experts_cpu(
1077 K,
1078 E,
1079 topk,
1060 - num_tokens_post_pad);
1080 + num_tokens_post_pad,
1081 + alpha.has_value() ? float(alpha.value()) : 0,
1082 + limit.has_value() ? float(limit.value()) : 0,
1083 + act_func,
1084 + with_bias);
1085 + } else if (moe_comp_method == CPUQuantMethod::MXFP4) {
1086 + scalar_t* __restrict__ A_tmp = (scalar_t*)((void*)(intermediate_cache2 + M * topk * K));
1087 + float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K));
1088 + scalar_t* __restrict__ intermediate_cache0 = (scalar_t*)((void*)(C_tmp + num_threads * 2 * BLOCK_M * BLOCK_N));
1089 + scalar_t* __restrict__ B_tmp = (scalar_t*)((void*)(intermediate_cache0 + M * topk * 2 * N));
1090 + bool with_bias = w1_bias.has_value();
1091 + auto act_func = alpha.has_value() && limit.has_value() ? CPUAcTMethod::swiglu : CPUAcTMethod::silu_and_mul;
1092 +
1093 + // mxfp4 supports only group size of 32 (2^5)
1094 + constexpr int64_t group_size = 32;
1095 + auto w1s = w1_scale.value();
1096 + auto w2s = w2_scale.value();
1097 + TORCH_CHECK(w1s.numel() == E * 2 * N * K / group_size, "w1_scale size mismatch");
1098 + TORCH_CHECK(w2s.numel() == E * K * N / group_size, "w2_scale size mismatch");
1099 + fused_experts_fp_kernel_impl<scalar_t, uint8_t, uint8_t, true>(
1100 + out_hidden_states.data_ptr<scalar_t>(),
1101 + intermediate_cache0,
1102 + intermediate_cache1,
1103 + intermediate_cache2,
1104 + A_tmp,
1105 + B_tmp,
1106 + C_tmp,
1107 + hidden_states.data_ptr<scalar_t>(),
1108 + packed_w1.data_ptr<uint8_t>(),
1109 + packed_w2.data_ptr<uint8_t>(),
1110 + with_bias ? w1_bias.value().data_ptr<float>() : nullptr,
1111 + with_bias ? w2_bias.value().data_ptr<float>() : nullptr,
1112 + w1s.data_ptr<uint8_t>(),
1113 + w2s.data_ptr<uint8_t>(),
1114 + /*block_size_N*/ 1,
1115 + /*block_size_K*/ group_size,
1116 + topk_weights_.data_ptr<float>(),
1117 + sorted_ids,
1118 + expert_ids,
1119 + offsets,
1120 + M,
1121 + N,
1122 + K,
1123 + E,
1124 + topk,
1125 + num_tokens_post_pad,
1126 + alpha.has_value() ? float(alpha.value()) : 0,
1127 + limit.has_value() ? float(limit.value()) : 0,
1128 + act_func,
1129 + with_bias);
1130 } else if (moe_comp_method == CPUQuantMethod::INT4_W4A8) {
1131 uint8_t* __restrict__ A_tmp = (uint8_t*)((void*)(intermediate_cache2 + M * topk * K));
1132 float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K));
@@ -1185,7 +1254,7 @@ at::Tensor shared_expert_cpu(
1254 CHECK_EQ(packed_w2.size(1), packed_N);
1255
1256 // check scales
1188 - check_moe_scales(use_int8_w8a8, use_fp8_w8a16, w1_scale, w2_scale, block_size);
1257 + check_moe_scales(use_int8_w8a8, use_fp8_w8a16, false, w1_scale, w2_scale, block_size);
1258
1259 at::Tensor out_hidden_states = inplace ? hidden_states : at::empty_like(hidden_states);
1260
csrc/cpu/sgl-kernels/moe.h
+106
@@ -176,3 +176,109 @@ inline void silu_and_mul_stub(
176 out_vec.store(out + d);
177 }
178 }
179 +
180 +template <typename scalar_t>
181 +inline void copy_mul_stub(scalar_t* __restrict__ out, const float* __restrict__ input, float weight, int64_t size) {
182 + using bVec = at::vec::Vectorized<scalar_t>;
183 + using fVec = at::vec::Vectorized<float>;
184 + constexpr int kVecSize = bVec::size();
185 + const fVec weight_vec = fVec(weight);
186 + int64_t d;
187 +#pragma GCC unroll 4
188 + for (d = 0; d <= size - kVecSize; d += kVecSize) {
189 + fVec data0 = fVec::loadu(input + d) * weight_vec;
190 + fVec data1 = fVec::loadu(input + d + fVec::size()) * weight_vec;
191 + bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
192 + out_vec.store(out + d);
193 + }
194 + for (; d < size; ++d) {
195 + out[d] = static_cast<scalar_t>(input[d] * weight);
196 + }
197 +}
198 +
199 +// input = input + input2
200 +inline void add_bias_stub(float* __restrict__ input, const float* __restrict__ input2, int64_t size) {
201 + using fVec = at::vec::Vectorized<float>;
202 + constexpr int kVecSize = fVec::size();
203 + int64_t d;
204 +#pragma GCC unroll 4
205 + for (d = 0; d <= size - kVecSize; d += kVecSize) {
206 + fVec x_fvec = fVec::loadu(input + d);
207 + fVec y_fvec = fVec::loadu(input2 + d);
208 + x_fvec = x_fvec + y_fvec;
209 + x_fvec.store(input + d);
210 + }
211 + for (; d < size; ++d) {
212 + input[d] = input[d] + input2[d];
213 + }
214 +}
215 +
216 +template <typename scalar_t>
217 +inline void copy_mul_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, float weight, int64_t size) {
218 + using bVec = at::vec::Vectorized<scalar_t>;
219 + using fVec = at::vec::Vectorized<float>;
220 + constexpr int kVecSize = bVec::size();
221 + const fVec weight_vec = fVec(weight);
222 + int64_t d;
223 +#pragma GCC unroll 4
224 + for (d = 0; d <= size - kVecSize; d += kVecSize) {
225 + bVec x = bVec::loadu(input + d);
226 + fVec x0, x1;
227 + std::tie(x0, x1) = at::vec::convert_to_float(x);
228 + x0 = x0 * weight_vec;
229 + x1 = x1 * weight_vec;
230 + bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
231 + out_vec.store(out + d);
232 + }
233 + for (; d < size; ++d) {
234 + out[d] = static_cast<scalar_t>(input[d] * weight);
235 + }
236 +}
237 +
238 +template <typename scalar_t>
239 +inline void clamp_sigmoid_and_mul_stub(
240 + scalar_t* __restrict__ out,
241 + const scalar_t* __restrict__ input,
242 + int64_t size,
243 + const float alpha,
244 + const float limit) {
245 + using bVec = at::vec::Vectorized<scalar_t>;
246 + using fVec = at::vec::Vectorized<float>;
247 + const fVec one = fVec(1.f);
248 + const fVec zero = fVec(0.f);
249 + const fVec limit_v = fVec(limit);
250 + const fVec nlimit_v = fVec(-limit);
251 + const fVec alpha_v = fVec(alpha);
252 +
253 + // no remainder
254 +#pragma GCC unroll 4
255 + for (int64_t d = 0; d < size; d += bVec::size()) {
256 + bVec x = bVec::loadu(input + d);
257 + fVec x0_, y0_;
258 + std::tie(x0_, y0_) = at::vec::convert_to_float(x);
259 + float tmp_buffer[fVec::size() * 2]; // 32
260 + float tmp_glu[fVec::size()]; // 16
261 + float tmp_linear[fVec::size()]; // 16
262 + x0_.store(tmp_buffer);
263 + y0_.store(tmp_buffer + fVec::size());
264 + // interleaved: x[2i] = glu, x[2i+1] = linear
265 + for (int j = 0; j < fVec::size(); ++j) {
266 + // x0 [0,2,..30]
267 + tmp_glu[j] = tmp_buffer[j * 2];
268 + // y0 [1,3,...31]
269 + tmp_linear[j] = tmp_buffer[j * 2 + 1];
270 + }
271 + fVec x0 = fVec::loadu(tmp_glu);
272 + fVec y0 = fVec::loadu(tmp_linear);
273 +
274 + // clamp
275 + x0 = at::vec::minimum(x0, limit_v);
276 + y0 = at::vec::minimum(limit_v, at::vec::maximum(nlimit_v, y0));
277 + // x * sigmoid(x * alpha)
278 + x0 = x0 / (one + (x0 * alpha_v).neg().exp_u20());
279 + // (y + 1) * x
280 + y0 = y0 + one;
281 + x0 = x0 * y0;
282 + convert_from_float_and_store<scalar_t>(out + d / 2, x0);
283 + }
284 +}
csrc/cpu/sgl-kernels/moe_fp8.cpp
+92 -56
@@ -7,8 +7,8 @@
7 #include "gemm.h"
8 #include "moe.h"
9
10 -template <typename scalar_t>
11 -void fused_experts_fp8_kernel_impl(
10 +template <typename scalar_t, typename packed_t, typename param_t, bool is_mxfp4>
11 +void fused_experts_fp_kernel_impl(
12 scalar_t* __restrict__ output,
13 scalar_t* __restrict__ ic0,
14 scalar_t* __restrict__ ic1,
@@ -17,10 +17,12 @@ void fused_experts_fp8_kernel_impl(
17 scalar_t* __restrict__ B_tmp,
18 float* __restrict__ C_tmp,
19 const scalar_t* __restrict__ input,
20 - const at::Float8_e4m3fn* __restrict__ packed_w1,
21 - const at::Float8_e4m3fn* __restrict__ packed_w2,
22 - const float* __restrict__ w1s,
23 - const float* __restrict__ w2s,
20 + const packed_t* __restrict__ packed_w1,
21 + const packed_t* __restrict__ packed_w2,
22 + const float* __restrict__ w1_bias,
23 + const float* __restrict__ w2_bias,
24 + const param_t* __restrict__ w1s,
25 + const param_t* __restrict__ w2s,
26 int64_t block_size_N,
27 int64_t block_size_K,
28 const float* __restrict__ topk_weights,
@@ -32,7 +34,11 @@ void fused_experts_fp8_kernel_impl(
34 int64_t K,
35 int64_t E,
36 int64_t topk,
35 - int64_t num_tokens_post_pad) {
37 + int64_t num_tokens_post_pad,
38 + float alpha,
39 + float limit,
40 + CPUAcTMethod act_func,
41 + bool with_bias) {
42 constexpr int64_t BLOCK_M = block_size_m();
43 constexpr int64_t BLOCK_N = block_size_n();
44
@@ -42,12 +48,20 @@ void fused_experts_fp8_kernel_impl(
48 int64_t scale_size_N = div_up(2 * N, block_size_N);
49 int64_t scale_size_K = div_up(K, block_size_K);
50 int64_t blocks_n_per_group = block_size_N / BLOCK_N;
51 + std::function<int64_t(int64_t)> scale_offset_per_block;
52 + if constexpr (is_mxfp4) {
53 + scale_offset_per_block = [&](int64_t a) { return a * BLOCK_N; };
54 + } else {
55 + scale_offset_per_block = [&](int64_t a) { return a / blocks_n_per_group; };
56 + }
57 +
58 + const int64_t packed_K = get_row_size<packed_t>(K);
59
46 - const int64_t stride_e = 2 * N * K;
47 - const int64_t stride_n = K;
60 + const int64_t stride_e = 2 * N * packed_K;
61 + const int64_t stride_n = packed_K;
62
63 int64_t avg_M = std::max(int64_t(1), M * topk / E);
50 - const bool use_brgemm = can_use_brgemm<at::Float8_e4m3fn>(avg_M);
64 + const bool use_brgemm = can_use_brgemm<packed_t>(avg_M);
65
66 int64_t B_tmp_size_per_thread = MAX_CACHE_BLOCK_SIZE * BLOCK_N * std::max(K, N);
67
@@ -57,14 +71,15 @@ void fused_experts_fp8_kernel_impl(
71 int tid = get_thread_num();
72 scalar_t* __restrict__ A = A_tmp + tid * BLOCK_M * K;
73
60 - loop_2d<at::Float8_e4m3fn>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
74 + loop_2d<packed_t>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
75 int64_t n_size = std::min(2 * N - nb * BLOCK_N, BLOCK_N);
76
77 // B shape [K, n_size] in vnni format
78 int32_t expert_id = expert_ids[mb];
65 - const at::Float8_e4m3fn* __restrict__ B = packed_w1 + expert_id * stride_e + nb * BLOCK_N * stride_n;
66 - const float* __restrict__ Bs =
67 - w1s + expert_id * scale_size_N * scale_size_K + (nb / blocks_n_per_group) * scale_size_K;
79 + const packed_t* __restrict__ B = packed_w1 + expert_id * stride_e + nb * BLOCK_N * stride_n;
80 + const param_t* __restrict__ Bs =
81 + w1s + expert_id * scale_size_N * scale_size_K + scale_offset_per_block(nb) * scale_size_K;
82 + const float* __restrict__ B_bias = with_bias ? w1_bias + expert_id * 2 * N + nb * BLOCK_N : nullptr;
83
84 // do unpacking for the first row or a new expert
85 int32_t pre_expert_id = mb == 0 ? -1 : expert_ids[mb - 1];
@@ -88,6 +103,7 @@ void fused_experts_fp8_kernel_impl(
103 /* C */ ic0 + offset * 2 * N + nb * BLOCK_N,
104 /* Btmp */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * K,
105 /* Ctmp */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
106 + /* Bbias */ B_bias,
107 /* scale */ Bs,
108 /* M */ m_size,
109 /* N */ n_size,
@@ -106,12 +122,20 @@ void fused_experts_fp8_kernel_impl(
122 });
123
124 // stage 1.5: intermediate_cache1 = silu(intermediate_cache0)
109 - at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
110 - for (int64_t m = begin; m < end; ++m) {
111 - silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N);
112 - }
113 - });
114 -
125 + if (act_func == CPUAcTMethod::silu_and_mul) {
126 + at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
127 + for (int64_t m = begin; m < end; ++m) {
128 + silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N);
129 + }
130 + });
131 + } else if (act_func == CPUAcTMethod::swiglu) {
132 + at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
133 + for (int64_t m = begin; m < end; ++m) {
134 + clamp_sigmoid_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, N, alpha, limit);
135 + clamp_sigmoid_and_mul_stub(ic1 + m * N + N / 2, ic0 + m * 2 * N + N, N, alpha, limit);
136 + }
137 + });
138 + }
139 // stage 2: intermediate_cache2 = intermediate_cache1 @ w2
140 // w2 : [E, K, N] as [E, OC, IC]
141 const int64_t OC = K; // rename K as OC
@@ -120,15 +144,16 @@ void fused_experts_fp8_kernel_impl(
144 const int64_t NB2 = div_up(OC, BLOCK_N);
145 scale_size_N = div_up(K, block_size_N);
146 scale_size_K = div_up(N, block_size_K);
123 - const int64_t stride_e2 = OC * IC;
124 - const int64_t stride_oc = IC;
147 + const int64_t packed_IC = get_row_size<packed_t>(IC);
148 + const int64_t stride_e2 = OC * packed_IC;
149 + const int64_t stride_oc = packed_IC;
150
151 // parallel on [MB2, NB2]
152 parallel_2d(MB2, NB2, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
153 int tid = get_thread_num();
154 alignas(64) scalar_t C[BLOCK_M * BLOCK_K];
155
131 - loop_2d<at::Float8_e4m3fn>(mb0, mb1, nb0, nb1, BLOCK_N * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
156 + loop_2d<packed_t>(mb0, mb1, nb0, nb1, BLOCK_N * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
157 int64_t m_size = offsets[mb + 1] - offsets[mb];
158 int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);
159
@@ -139,9 +164,10 @@ void fused_experts_fp8_kernel_impl(
164
165 // B shape [IC, n_size] in vnni format
166 int32_t expert_id = expert_ids[mb];
142 - const at::Float8_e4m3fn* __restrict__ B = packed_w2 + expert_id * stride_e2 + nb * BLOCK_N * stride_oc;
143 - const float* __restrict__ Bs =
144 - w2s + expert_id * scale_size_N * scale_size_K + (nb / blocks_n_per_group) * scale_size_K;
167 + const packed_t* __restrict__ B = packed_w2 + expert_id * stride_e2 + nb * BLOCK_N * stride_oc;
168 + const param_t* __restrict__ Bs =
169 + w2s + expert_id * scale_size_N * scale_size_K + scale_offset_per_block(nb) * scale_size_K;
170 + const float* __restrict__ B_bias = with_bias ? w2_bias + expert_id * OC + nb * BLOCK_N : nullptr;
171
172 // do unpacking for the first row or a new expert
173 int32_t pre_expert_id = mb == 0 ? -1 : expert_ids[mb - 1];
@@ -153,6 +179,7 @@ void fused_experts_fp8_kernel_impl(
179 /* C */ C,
180 /* Btmp */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * IC,
181 /* Ctmp */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
182 + /* Bbias */ B_bias,
183 /* scale */ Bs,
184 /* M */ m_size,
185 /* N */ n_size,
@@ -177,7 +204,6 @@ void fused_experts_fp8_kernel_impl(
204 at::native::cpublas::brgemm_release();
205 }
206 });
180 -
207 // stage 3: out = intermediate_cache2.sum(dim=1)
208 // from [M, topk, K] to [M, K]
209 at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
@@ -187,35 +213,43 @@ void fused_experts_fp8_kernel_impl(
213 });
214 }
215
190 -#define INSTANTIATE_MOE_FP8_TEMPLATE(TYPE) \
191 - template void fused_experts_fp8_kernel_impl<TYPE>( \
192 - TYPE* __restrict__ output, \
193 - TYPE* __restrict__ ic0, \
194 - TYPE* __restrict__ ic1, \
195 - TYPE* __restrict__ ic2, \
196 - TYPE* __restrict__ A_tmp, \
197 - TYPE* __restrict__ B_tmp, \
198 - float* __restrict__ C_tmp, \
199 - const TYPE* __restrict__ input, \
200 - const at::Float8_e4m3fn* __restrict__ packed_w1, \
201 - const at::Float8_e4m3fn* __restrict__ packed_w2, \
202 - const float* __restrict__ w1s, \
203 - const float* __restrict__ w2s, \
204 - int64_t block_size_N, \
205 - int64_t block_size_K, \
206 - const float* __restrict__ topk_weights, \
207 - const int32_t* __restrict__ sorted_ids, \
208 - const int32_t* __restrict__ expert_ids, \
209 - const int32_t* __restrict__ offsets, \
210 - int64_t M, \
211 - int64_t N, \
212 - int64_t K, \
213 - int64_t E, \
214 - int64_t topk, \
215 - int64_t num_tokens_post_pad)
216 -
217 -INSTANTIATE_MOE_FP8_TEMPLATE(at::BFloat16);
218 -INSTANTIATE_MOE_FP8_TEMPLATE(at::Half);
216 +#define INSTANTIATE_MOE_FP_TEMPLATE(TYPE1, TYPE2, TYPE3, IS_MXFP4) \
217 + template void fused_experts_fp_kernel_impl<TYPE1, TYPE2, TYPE3, IS_MXFP4>( \
218 + TYPE1* __restrict__ output, \
219 + TYPE1* __restrict__ ic0, \
220 + TYPE1* __restrict__ ic1, \
221 + TYPE1* __restrict__ ic2, \
222 + TYPE1* __restrict__ A_tmp, \
223 + TYPE1* __restrict__ B_tmp, \
224 + float* __restrict__ C_tmp, \
225 + const TYPE1* __restrict__ input, \
226 + const TYPE2* __restrict__ packed_w1, \
227 + const TYPE2* __restrict__ packed_w2, \
228 + const float* __restrict__ w1_bias, \
229 + const float* __restrict__ w2_bias, \
230 + const TYPE3* __restrict__ w1s, \
231 + const TYPE3* __restrict__ w2s, \
232 + int64_t block_size_N, \
233 + int64_t block_size_K, \
234 + const float* __restrict__ topk_weights, \
235 + const int32_t* __restrict__ sorted_ids, \
236 + const int32_t* __restrict__ expert_ids, \
237 + const int32_t* __restrict__ offsets, \
238 + int64_t M, \
239 + int64_t N, \
240 + int64_t K, \
241 + int64_t E, \
242 + int64_t topk, \
243 + int64_t num_tokens_post_pad, \
244 + float alpha, \
245 + float limit, \
246 + CPUAcTMethod act_func, \
247 + bool with_bias)
248 +
249 +INSTANTIATE_MOE_FP_TEMPLATE(at::BFloat16, at::Float8_e4m3fn, float, false);
250 +INSTANTIATE_MOE_FP_TEMPLATE(at::Half, at::Float8_e4m3fn, float, false);
251 +INSTANTIATE_MOE_FP_TEMPLATE(at::BFloat16, uint8_t, uint8_t, true);
252 +INSTANTIATE_MOE_FP_TEMPLATE(at::Half, uint8_t, uint8_t, true);
253
254 template <typename scalar_t>
255 void shared_expert_fp8_kernel_impl(
@@ -266,6 +300,7 @@ void shared_expert_fp8_kernel_impl(
300 /* C */ ic0 + mb * BLOCK_M * 2 * N + nb * BLOCK_N,
301 /* Btmp */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * K,
302 /* Ctmp */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
303 + /* Bbias */ nullptr,
304 /* scale */ w1s + (nb / blocks_n_per_group) * scale_size_K,
305 /* M */ m_size,
306 /* N */ n_size,
@@ -317,6 +352,7 @@ void shared_expert_fp8_kernel_impl(
352 /* C */ C,
353 /* Btmp */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * IC,
354 /* Ctmp */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
355 + /* Bbias */ nullptr,
356 /* scale */ w2s + (nb / blocks_n_per_group) * scale_size_K,
357 /* M */ m_size,
358 /* N */ n_size,
csrc/cpu/sgl-kernels/vec.h
+15 -2
@@ -11,7 +11,7 @@
11
12 #include <ATen/cpu/vec/functional.h>
13 #include <ATen/cpu/vec/vec.h>
14 -
14 +#include <immintrin.h>
15 namespace {
16
17 using namespace at::vec;
@@ -21,6 +21,15 @@ inline Vectorized<scalar_t> convert_from_float_ext(const Vectorized<float>& a, c
21 return at::vec::convert_from_float<scalar_t>(a, b);
22 }
23
24 +template <typename scalar_t>
25 +inline void convert_from_float_and_store(scalar_t* out, const Vectorized<float>& a) {
26 + float out_buffer[at::vec::Vectorized<float>::size()];
27 + a.store(out_buffer);
28 + for (int i = 0; i < 16; i++) {
29 + out[i] = (scalar_t)out_buffer[i];
30 + }
31 +}
32 +
33 // allow f16, bf16
34 template <typename scalar_t, typename std::enable_if_t<is_reduced_floating_point_v<scalar_t>, int> = 1>
35 inline std::tuple<Vectorized<float>, Vectorized<float>> load_float_vec2(const scalar_t* __restrict__ data) {
@@ -50,6 +59,11 @@ convert_from_float_ext<at::BFloat16>(const Vectorized<float>& a, const Vectorize
59 return (__m512i)(_mm512_cvtne2ps_pbh(__m512(b), __m512(a)));
60 }
61
62 +template <>
63 +inline void convert_from_float_and_store<at::BFloat16>(at::BFloat16* out, const Vectorized<float>& a) {
64 + _mm256_storeu_si256((__m256i*)out, (__m256i)(_mm512_cvtneps_pbh(__m512(a))));
65 +}
66 +
67 #define CVT_BF16_TO_FP32(a) _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(a), 16))
68
69 #define CVT_FP16_TO_FP32(a) _mm512_cvtph_ps(a)
@@ -125,7 +139,6 @@ inline __m512bh CVT_FP8_TO_BF16(__m256i a) {
139 return cvt_e4m3_bf16_intrinsic_with_denorm(a);
140 #endif
141 }
128 -
142 // faster version of float8_e4m3fn conversion to bfloat16
143 //
144 // we mapped cuda implementation from below link and vectorized with avx512:
csrc/cpu/torch_bindings.cpp
+12 -2
@@ -64,6 +64,8 @@ at::Tensor weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2,
64
65 at::Tensor convert_weight_packed(at::Tensor& weight);
66
67 +at::Tensor convert_scale_packed(at::Tensor& scale);
68 +
69 at::Tensor fused_experts_cpu(
70 at::Tensor& hidden_states, at::Tensor& w1, at::Tensor& w2,
71 at::Tensor& topk_weights, at::Tensor& topk_ids, bool inplace,
@@ -71,7 +73,11 @@ at::Tensor fused_experts_cpu(
73 const std::optional<at::Tensor>& w2_scale,
74 const std::optional<at::Tensor>& w1_zero,
75 const std::optional<at::Tensor>& w2_zero,
74 - const std::optional<std::vector<int64_t>> block_size, bool is_vnni);
76 + const std::optional<std::vector<int64_t>> block_size,
77 + const std::optional<at::Tensor>& w1_bias,
78 + const std::optional<at::Tensor>& w2_bias,
79 + const std::optional<double>& alpha, const std::optional<double>& limit,
80 + bool is_vnni);
81
82 at::Tensor int8_scaled_mm_with_quant(at::Tensor& mat1, at::Tensor& mat2,
83 at::Tensor& scales2,
@@ -404,11 +410,15 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
410 ops.impl("weight_packed_linear", torch::kCPU, &weight_packed_linear);
411 ops.def("convert_weight_packed(Tensor! weight) -> Tensor");
412 ops.impl("convert_weight_packed", torch::kCPU, &convert_weight_packed);
413 + ops.def("convert_scale_packed(Tensor! scale) -> Tensor");
414 + ops.impl("convert_scale_packed", torch::kCPU, &convert_scale_packed);
415 ops.def(
416 "fused_experts_cpu(Tensor hidden_states, Tensor w1, Tensor w2, Tensor "
417 "topk_weights, Tensor topk_ids, bool "
418 "inplace, int moe_comp_method, Tensor? w1_scale, Tensor? w2_scale, "
411 - "Tensor? w1_zero, Tensor? w2_zero, int[]? block_size, bool is_vnni) -> "
419 + "Tensor? w1_zero, Tensor? w2_zero, int[]? block_size, "
420 + "Tensor? w1_bias, Tensor? w2_bias, float? alpha, float? limit, "
421 + "bool is_vnni) -> "
422 "Tensor");
423 ops.impl("fused_experts_cpu", torch::kCPU, &fused_experts_cpu);
424 ops.def(
tests/kernels/moe/test_cpu_fp8_fused_moe.py deleted
-256
@@ -1,256 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Tests for CPU FP8 W8A16 block-scaled fused MoE kernel."""
4 -
5 -import math
6 -import sys
7 -
8 -import pytest
9 -import torch
10 -import torch.nn.functional as F
11 -
12 -from vllm.platforms import current_platform
13 -from vllm.utils.torch_utils import set_random_seed
14 -
15 -if not current_platform.is_cpu():
16 - pytest.skip("skipping CPU-only tests", allow_module_level=True)
17 -
18 -import vllm._custom_ops as ops # noqa: E402
19 -
20 -if not hasattr(torch.ops._C, "fused_experts_cpu"):
21 - pytest.skip("fused_experts_cpu op not available", allow_module_level=True)
22 -
23 -
24 -BLOCK_SIZE = [128, 128] # [block_n, block_k]
25 -
26 -# FP8 weight generation parameters
27 -_FP8_INFO = torch.finfo(torch.float8_e4m3fn)
28 -FP8_SCALE = _FP8_INFO.max # 448.0
29 -FACTOR_FOR_SCALE = 1e-3
30 -
31 -# Tolerance for FP8 W8A16
32 -FP8_W8A16_ATOL = 1e-2
33 -FP8_W8A16_RTOL = 1e-2
34 -
35 -
36 -def _silu_and_mul(x: torch.Tensor) -> torch.Tensor:
37 - d = x.shape[-1] // 2
38 - return F.silu(x[..., :d]) * x[..., d:]
39 -
40 -
41 -def _block_dequant_weight(
42 - weight: torch.Tensor,
43 - scales: torch.Tensor,
44 - block_size: list[int],
45 -) -> torch.Tensor:
46 - """Block-dequantize FP8 weight [E, N, K] → float [E, N, K].
47 -
48 - Each (block_n × block_k) tile is multiplied by its per-block scale.
49 - """
50 - E, N, K = weight.shape
51 - block_n, block_k = block_size
52 - pad_N = (block_n - N % block_n) % block_n
53 - pad_K = (block_k - K % block_k) % block_k
54 -
55 - if pad_N > 0 or pad_K > 0:
56 - weight = F.pad(weight, (0, pad_K, 0, pad_N))
57 -
58 - n_tiles = math.ceil(N / block_n)
59 - k_tiles = math.ceil(K / block_k)
60 -
61 - weight_block = (
62 - weight.view(E, n_tiles, block_n, k_tiles, block_k)
63 - .permute(0, 1, 3, 2, 4)
64 - .float()
65 - .contiguous()
66 - )
67 - weight_scaled = (
68 - (weight_block * scales.view(E, n_tiles, k_tiles, 1, 1))
69 - .permute(0, 1, 3, 2, 4)
70 - .contiguous()
71 - )
72 - if pad_N > 0 or pad_K > 0:
73 - weight_scaled = weight_scaled.view(E, N + pad_N, K + pad_K)
74 - weight_scaled = weight_scaled[..., :N, :K].contiguous()
75 - else:
76 - weight_scaled = weight_scaled.view(E, N, K)
77 - return weight_scaled
78 -
79 -
80 -def ref_w8a16_block_fp8_moe(
81 - a: torch.Tensor,
82 - w1: torch.Tensor,
83 - w2: torch.Tensor,
84 - w1_s: torch.Tensor,
85 - w2_s: torch.Tensor,
86 - topk_weight: torch.Tensor,
87 - topk_ids: torch.Tensor,
88 - block_size: list[int],
89 -) -> torch.Tensor:
90 - """Reference FP8 W8A16 block-scaled fused MoE in pure torch.
91 -
92 - Steps:
93 - 1. Block-dequant FP8 weights → float
94 - 2. For each expert: matmul → SiLU+Mul → matmul
95 - 3. Weighted sum across top-k experts
96 - """
97 - B, D = a.shape
98 - topk = topk_ids.size(1)
99 -
100 - w1_dq = _block_dequant_weight(w1, w1_s, block_size)
101 - w2_dq = _block_dequant_weight(w2, w2_s, block_size)
102 -
103 - a_exp = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D).float()
104 - out = torch.zeros(B * topk, w2_dq.shape[1], dtype=torch.float32)
105 -
106 - topk_weight_flat = topk_weight.view(-1)
107 - topk_ids_flat = topk_ids.view(-1)
108 -
109 - for i in range(w1_dq.shape[0]):
110 - mask = topk_ids_flat == i
111 - if mask.sum():
112 - ic0 = torch.matmul(a_exp[mask], w1_dq[i].transpose(0, 1))
113 - ic1 = _silu_and_mul(ic0)
114 - out[mask] = torch.matmul(ic1, w2_dq[i].transpose(0, 1))
115 -
116 - return (
117 - (out.view(B, -1, w2_dq.shape[1]) * topk_weight_flat.view(B, -1, 1))
118 - .sum(dim=1)
119 - .to(a.dtype)
120 - )
121 -
122 -
123 -def _make_fp8_moe_weights(
124 - E: int,
125 - N: int,
126 - K: int,
127 - block_size: list[int],
128 -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
129 - """Generate random FP8 MoE weights with random block scales.
130 -
131 - Weight generation follows SGLang: ``randn * FP8_SCALE`` → clamp → cast.
132 - Scales are small random values (``FACTOR_FOR_SCALE``), independent of
133 - the actual weight magnitudes — this is sufficient to test the kernel's
134 - block-dequant + matmul correctness.
135 -
136 - Returns: (w1, w2, w1_s, w2_s)
137 - """
138 - block_n, block_k = block_size
139 -
140 - w1 = (
141 - (torch.randn(E, 2 * N, K) * FP8_SCALE)
142 - .clamp(min=-FP8_SCALE, max=FP8_SCALE)
143 - .to(torch.float8_e4m3fn)
144 - )
145 - w2 = (
146 - (torch.randn(E, K, N) * FP8_SCALE)
147 - .clamp(min=-FP8_SCALE, max=FP8_SCALE)
148 - .to(torch.float8_e4m3fn)
149 - )
150 -
151 - w1_s = (
152 - torch.randn(E, math.ceil(2 * N / block_n), math.ceil(K / block_k))
153 - * FACTOR_FOR_SCALE
154 - )
155 - w2_s = (
156 - torch.randn(E, math.ceil(K / block_n), math.ceil(N / block_k))
157 - * FACTOR_FOR_SCALE
158 - )
159 - return w1, w2, w1_s, w2_s
160 -
161 -
162 -def _prepack_experts(w: torch.Tensor) -> torch.Tensor:
163 - """VNNI-prepack each expert's weight via ``convert_weight_packed``."""
164 - return torch.stack(
165 - [torch.ops._C.convert_weight_packed(w[e]) for e in range(w.shape[0])]
166 - )
167 -
168 -
169 -NUM_TOKENS = [1, 2, 64, 121]
170 -# (M, intermediate_size N, hidden_size K, num_experts E, topk)
171 -MoE_CONFIGS = [
172 - (256, 512, 8, 2),
173 - (256, 512, 8, 4),
174 - (512, 256, 8, 2),
175 - (512, 256, 8, 4),
176 - (512, 512, 8, 2),
177 - (512, 512, 8, 4),
178 - (768, 2048, 8, 2),
179 - (768, 2048, 8, 4),
180 - (768, 2048, 128, 8),
181 -]
182 -SEEDS = [0]
183 -
184 -
185 -@pytest.mark.parametrize("M", NUM_TOKENS)
186 -@pytest.mark.parametrize("N,K,E,topk", MoE_CONFIGS)
187 -@pytest.mark.parametrize("seed", SEEDS)
188 -def test_w8a16_block_fp8_cpu_fused_moe(M, N, K, E, topk, seed):
189 - """Test fused_experts_cpu FP8 W8A16 against dequantised torch reference."""
190 - set_random_seed(seed)
191 -
192 - a = torch.randn(M, K, dtype=torch.bfloat16) / math.sqrt(K)
193 - w1, w2, w1_s, w2_s = _make_fp8_moe_weights(E, N, K, BLOCK_SIZE)
194 -
195 - score = torch.randn(M, E, dtype=torch.bfloat16)
196 - score = torch.softmax(score, dim=-1, dtype=torch.float32)
197 - topk_weight, topk_ids = torch.topk(score, topk)
198 - topk_ids = topk_ids.to(torch.int32)
199 -
200 - ref_out = ref_w8a16_block_fp8_moe(
201 - a,
202 - w1,
203 - w2,
204 - w1_s,
205 - w2_s,
206 - topk_weight,
207 - topk_ids,
208 - BLOCK_SIZE,
209 - )
210 -
211 - pw1, pw2 = _prepack_experts(w1), _prepack_experts(w2)
212 -
213 - # Test inplace=False against reference
214 - out = ops.fused_experts_cpu(
215 - a.clone(),
216 - pw1,
217 - pw2,
218 - topk_weight,
219 - topk_ids,
220 - False,
221 - ops.CPUQuantMethod.FP8_W8A16,
222 - w1_s,
223 - w2_s,
224 - None,
225 - None,
226 - BLOCK_SIZE,
227 - True,
228 - )
229 - torch.testing.assert_close(
230 - ref_out.bfloat16(),
231 - out,
232 - atol=FP8_W8A16_ATOL,
233 - rtol=FP8_W8A16_RTOL,
234 - )
235 -
236 - # Test inplace=True produces identical output
237 - out_inplace = ops.fused_experts_cpu(
238 - a.clone(),
239 - pw1,
240 - pw2,
241 - topk_weight,
242 - topk_ids,
243 - True,
244 - ops.CPUQuantMethod.FP8_W8A16,
245 - w1_s,
246 - w2_s,
247 - None,
248 - None,
249 - BLOCK_SIZE,
250 - True,
251 - )
252 - torch.testing.assert_close(out_inplace, out, atol=0, rtol=0)
253 -
254 -
255 -if __name__ == "__main__":
256 - sys.exit(pytest.main([__file__, "-v"]))
tests/kernels/moe/test_cpu_quant_fused_moe.py new
+500
@@ -0,0 +1,500 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for CPU quantized fused MoE kernels (FP8 W8A16 and MXFP4 W4A16)."""
4 +
5 +import math
6 +import sys
7 +
8 +import pytest
9 +import torch
10 +import torch.nn.functional as F
11 +
12 +from vllm.platforms import current_platform
13 +from vllm.utils.torch_utils import set_random_seed
14 +
15 +if not current_platform.is_cpu():
16 + pytest.skip("skipping CPU-only tests", allow_module_level=True)
17 +
18 +import vllm._custom_ops as ops # noqa: E402
19 +
20 +if not hasattr(torch.ops._C, "fused_experts_cpu"):
21 + pytest.skip("fused_experts_cpu op not available", allow_module_level=True)
22 +
23 +
24 +def _silu_and_mul(x: torch.Tensor) -> torch.Tensor:
25 + d = x.shape[-1] // 2
26 + return F.silu(x[..., :d]) * x[..., d:]
27 +
28 +
29 +def _prepack_experts(w: torch.Tensor) -> torch.Tensor:
30 + """VNNI-prepack expert weights via ``convert_weight_packed``."""
31 + return torch.ops._C.convert_weight_packed(w)
32 +
33 +
34 +# FP8 W8A16 block-scaled fused MoE
35 +
36 +BLOCK_SIZE = [128, 128] # [block_n, block_k]
37 +
38 +_FP8_INFO = torch.finfo(torch.float8_e4m3fn)
39 +FP8_SCALE = _FP8_INFO.max # 448.0
40 +FACTOR_FOR_SCALE = 1e-3
41 +
42 +
43 +def _block_dequant_weight(
44 + weight: torch.Tensor,
45 + scales: torch.Tensor,
46 + block_size: list[int],
47 +) -> torch.Tensor:
48 + """Block-dequantize FP8 weight [E, N, K] -> float [E, N, K]."""
49 + E, N, K = weight.shape
50 + block_n, block_k = block_size
51 + pad_N = (block_n - N % block_n) % block_n
52 + pad_K = (block_k - K % block_k) % block_k
53 +
54 + if pad_N > 0 or pad_K > 0:
55 + weight = F.pad(weight, (0, pad_K, 0, pad_N))
56 +
57 + n_tiles = math.ceil(N / block_n)
58 + k_tiles = math.ceil(K / block_k)
59 +
60 + weight_block = (
61 + weight.view(E, n_tiles, block_n, k_tiles, block_k)
62 + .permute(0, 1, 3, 2, 4)
63 + .float()
64 + .contiguous()
65 + )
66 + weight_scaled = (
67 + (weight_block * scales.view(E, n_tiles, k_tiles, 1, 1))
68 + .permute(0, 1, 3, 2, 4)
69 + .contiguous()
70 + )
71 + if pad_N > 0 or pad_K > 0:
72 + weight_scaled = weight_scaled.view(E, N + pad_N, K + pad_K)
73 + weight_scaled = weight_scaled[..., :N, :K].contiguous()
74 + else:
75 + weight_scaled = weight_scaled.view(E, N, K)
76 + return weight_scaled
77 +
78 +
79 +def ref_w8a16_block_fp8_moe(
80 + a: torch.Tensor,
81 + w1: torch.Tensor,
82 + w2: torch.Tensor,
83 + w1_s: torch.Tensor,
84 + w2_s: torch.Tensor,
85 + topk_weight: torch.Tensor,
86 + topk_ids: torch.Tensor,
87 + block_size: list[int],
88 +) -> torch.Tensor:
89 + """Reference FP8 W8A16 block-scaled fused MoE in pure torch."""
90 + B, D = a.shape
91 + topk = topk_ids.size(1)
92 +
93 + w1_dq = _block_dequant_weight(w1, w1_s, block_size)
94 + w2_dq = _block_dequant_weight(w2, w2_s, block_size)
95 +
96 + a_exp = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D).float()
97 + out = torch.zeros(B * topk, w2_dq.shape[1], dtype=torch.float32)
98 +
99 + topk_weight_flat = topk_weight.view(-1)
100 + topk_ids_flat = topk_ids.view(-1)
101 +
102 + for i in range(w1_dq.shape[0]):
103 + mask = topk_ids_flat == i
104 + if mask.sum():
105 + ic0 = torch.matmul(a_exp[mask], w1_dq[i].transpose(0, 1))
106 + ic1 = _silu_and_mul(ic0)
107 + out[mask] = torch.matmul(ic1, w2_dq[i].transpose(0, 1))
108 +
109 + return (
110 + (out.view(B, -1, w2_dq.shape[1]) * topk_weight_flat.view(B, -1, 1))
111 + .sum(dim=1)
112 + .to(a.dtype)
113 + )
114 +
115 +
116 +def _make_fp8_moe_weights(
117 + E: int,
118 + N: int,
119 + K: int,
120 + block_size: list[int],
121 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
122 + """Generate random FP8 MoE weights with random block scales."""
123 + block_n, block_k = block_size
124 +
125 + w1 = (
126 + (torch.randn(E, 2 * N, K) * FP8_SCALE)
127 + .clamp(min=-FP8_SCALE, max=FP8_SCALE)
128 + .to(torch.float8_e4m3fn)
129 + )
130 + w2 = (
131 + (torch.randn(E, K, N) * FP8_SCALE)
132 + .clamp(min=-FP8_SCALE, max=FP8_SCALE)
133 + .to(torch.float8_e4m3fn)
134 + )
135 +
136 + w1_s = (
137 + torch.randn(E, math.ceil(2 * N / block_n), math.ceil(K / block_k))
138 + * FACTOR_FOR_SCALE
139 + )
140 + w2_s = (
141 + torch.randn(E, math.ceil(K / block_n), math.ceil(N / block_k))
142 + * FACTOR_FOR_SCALE
143 + )
144 + return w1, w2, w1_s, w2_s
145 +
146 +
147 +FP8_NUM_TOKENS = [1, 2, 64, 121]
148 +FP8_MOE_CONFIGS = [
149 + (256, 512, 8, 2),
150 + (256, 512, 8, 4),
151 + (512, 256, 8, 2),
152 + (512, 256, 8, 4),
153 + (512, 512, 8, 2),
154 + (512, 512, 8, 4),
155 + (768, 2048, 8, 2),
156 + (768, 2048, 8, 4),
157 + (768, 2048, 128, 8),
158 +]
159 +
160 +
161 +@pytest.mark.parametrize("M", FP8_NUM_TOKENS)
162 +@pytest.mark.parametrize("N,K,E,topk", FP8_MOE_CONFIGS)
163 +@pytest.mark.parametrize("seed", [0])
164 +def test_w8a16_block_fp8_cpu_fused_moe(M, N, K, E, topk, seed):
165 + """Test fused_experts_cpu FP8 W8A16 against dequantised torch reference."""
166 + set_random_seed(seed)
167 +
168 + a = torch.randn(M, K, dtype=torch.bfloat16) / math.sqrt(K)
169 + w1, w2, w1_s, w2_s = _make_fp8_moe_weights(E, N, K, BLOCK_SIZE)
170 +
171 + score = torch.randn(M, E, dtype=torch.bfloat16)
172 + score = torch.softmax(score, dim=-1, dtype=torch.float32)
173 + topk_weight, topk_ids = torch.topk(score, topk)
174 + topk_ids = topk_ids.to(torch.int32)
175 +
176 + ref_out = ref_w8a16_block_fp8_moe(
177 + a, w1, w2, w1_s, w2_s, topk_weight, topk_ids, BLOCK_SIZE
178 + )
179 +
180 + pw1, pw2 = _prepack_experts(w1), _prepack_experts(w2)
181 +
182 + # Test inplace=False against reference
183 + out = ops.fused_experts_cpu(
184 + a.clone(),
185 + pw1,
186 + pw2,
187 + topk_weight,
188 + topk_ids,
189 + False,
190 + ops.CPUQuantMethod.FP8_W8A16,
191 + w1_s,
192 + w2_s,
193 + None,
194 + None,
195 + BLOCK_SIZE,
196 + is_vnni=True,
197 + )
198 + torch.testing.assert_close(ref_out.bfloat16(), out, atol=1e-2, rtol=1e-2)
199 +
200 + # Test inplace=True produces identical output
201 + out_inplace = ops.fused_experts_cpu(
202 + a.clone(),
203 + pw1,
204 + pw2,
205 + topk_weight,
206 + topk_ids,
207 + True,
208 + ops.CPUQuantMethod.FP8_W8A16,
209 + w1_s,
210 + w2_s,
211 + None,
212 + None,
213 + BLOCK_SIZE,
214 + is_vnni=True,
215 + )
216 + torch.testing.assert_close(out_inplace, out, atol=0, rtol=0)
217 +
218 +
219 +# MXFP4 W4A16 fused MoE
220 +
221 +
222 +class MXFP4QuantizeUtil:
223 + """MXFP4 quantization utility."""
224 +
225 + E2M1_max = 6.0
226 + E2M1_values = [0, 0.5, 1, 1.5, 2, 3, 4, 6]
227 + E2M1_bounds = torch.tensor([0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5])
228 + block_size = 32
229 +
230 + @classmethod
231 + def quantize(cls, input: torch.Tensor) -> tuple:
232 + """Quantize BF16 tensor to MXFP4 packed uint8 format."""
233 +
234 + def cast_fp4(x):
235 + sign = torch.sign(x)
236 + sign_bit = (2 - sign) // 2
237 + ord_ = torch.sum(
238 + (x.abs().unsqueeze(-1) - cls.E2M1_bounds.to(x.device)) > 0, dim=-1
239 + )
240 + fp4_val = (sign_bit * 0b1000 + ord_).to(torch.uint8)
241 + return fp4_val
242 +
243 + def fuse_uint4_to_uint8(x):
244 + left_side = x[..., 0::2]
245 + right_side = x[..., 1::2]
246 + new_data = right_side.clone() << 4
247 + new_data[..., : left_side.shape[-1]] += left_side
248 + return new_data
249 +
250 + original_shape = input.shape
251 + input = input.view(-1, cls.block_size)
252 + input_amax = input.abs().max(dim=-1, keepdim=True).values
253 + descale = input_amax / cls.E2M1_max
254 + min_value = torch.tensor(-127.0, device=descale.device)
255 + e8m0_scale = torch.ceil(torch.maximum(torch.log2(descale), min_value))
256 +
257 + input = (input / torch.exp2(e8m0_scale)).view(original_shape)
258 + input_q = cast_fp4(input)
259 + input_q = fuse_uint4_to_uint8(input_q)
260 + e8m0_scale = (e8m0_scale + 127).to(torch.uint8)
261 + return input_q, e8m0_scale
262 +
263 + @classmethod
264 + def dequantize(cls, quantized_data, dtype: torch.dtype, scale):
265 + """Dequantize MXFP4 packed tensor back to float."""
266 +
267 + def unfuse_uint8_to_uint4(x):
268 + left_side = x & 0x0F
269 + right_side = (x >> 4) & 0x0F
270 + shape = list(x.shape)
271 + shape[-1] = shape[-1] * 2
272 + result = torch.zeros(shape, dtype=torch.uint8, device=x.device)
273 + result[..., 0::2] = left_side
274 + result[..., 1::2] = right_side
275 + return result
276 +
277 + e8m0_scale = scale
278 + x_unfused = unfuse_uint8_to_uint4(quantized_data)
279 + sign = 1 - 2 * ((x_unfused & 0b1000) >> 3).to(torch.float32)
280 + magnitude = (x_unfused & 0b0111).to(torch.long)
281 + values = torch.tensor(cls.E2M1_values, device=quantized_data.device)
282 + original_shape = magnitude.shape
283 + x_float = values[magnitude.reshape(-1)].reshape(original_shape)
284 + x_float = sign.float() * x_float
285 + x_float = x_float.reshape(-1, cls.block_size)
286 + scale_factor = torch.exp2(e8m0_scale.float() - 127)
287 + scale_factor = scale_factor.reshape(-1, 1)
288 + x_float = x_float * scale_factor
289 + return x_float.reshape(original_shape).to(dtype)
290 +
291 +
292 +def _swiglu(x: torch.Tensor, alpha: float, limit: float) -> torch.Tensor:
293 + """SwigLU activation used in GPT-OSS.
294 + Input is interleaved: [gate_0, up_0, gate_1, up_1, ...] in last dim.
295 + """
296 + gate = x[..., 0::2]
297 + up = x[..., 1::2]
298 + gate_clamped = torch.clamp(gate, max=limit)
299 + up_clamped = torch.clamp(up, min=-limit, max=limit)
300 + return gate_clamped * torch.sigmoid(alpha * gate_clamped) * (up_clamped + 1)
301 +
302 +
303 +def ref_mxfp4_fused_moe(
304 + a: torch.Tensor,
305 + w1_dq: torch.Tensor,
306 + w2_dq: torch.Tensor,
307 + topk_weight: torch.Tensor,
308 + topk_ids: torch.Tensor,
309 + topk: int,
310 +) -> torch.Tensor:
311 + """Reference MXFP4 fused MoE with SiLU activation."""
312 + B, D = a.shape
313 + a_f = a.float()
314 + out = torch.zeros(B * topk, w2_dq.shape[1], dtype=torch.float32)
315 + topk_ids_flat = topk_ids.view(-1)
316 +
317 + for i in range(w1_dq.shape[0]):
318 + mask = topk_ids_flat == i
319 + if mask.sum() == 0:
320 + continue
321 + token_indices = torch.where(mask)[0]
322 + source_indices = token_indices // topk
323 + ic0 = torch.matmul(a_f[source_indices], w1_dq[i].float().T)
324 + ic1 = _silu_and_mul(ic0)
325 + out[mask] = torch.matmul(ic1, w2_dq[i].float().T)
326 +
327 + return (out.view(B, topk, -1) * topk_weight.unsqueeze(-1)).sum(dim=1).to(a.dtype)
328 +
329 +
330 +def ref_mxfp4_fused_moe_gptoss(
331 + a: torch.Tensor,
332 + w1_dq: torch.Tensor,
333 + w2_dq: torch.Tensor,
334 + w1_bias: torch.Tensor,
335 + w2_bias: torch.Tensor,
336 + topk_weight: torch.Tensor,
337 + topk_ids: torch.Tensor,
338 + alpha: float,
339 + limit: float,
340 +) -> torch.Tensor:
341 + """Reference MXFP4 fused MoE with SwigLU+bias (GPT-OSS style)."""
342 + B, D = a.shape
343 + topk = topk_ids.shape[1]
344 + a_f = a.float()
345 + E = w1_dq.shape[0]
346 + out = torch.zeros(B * topk, w2_dq.shape[1], dtype=torch.float32)
347 + topk_ids_flat = topk_ids.view(-1)
348 +
349 + for i in range(E):
350 + mask = topk_ids_flat == i
351 + if mask.sum() == 0:
352 + continue
353 + token_indices = torch.where(mask)[0]
354 + source_indices = token_indices // topk
355 + ic0 = torch.matmul(a_f[source_indices], w1_dq[i].float().T)
356 + ic0 = ic0 + w1_bias[i].float()
357 + ic1 = _swiglu(ic0, alpha, limit)
358 + ic2 = torch.matmul(ic1, w2_dq[i].float().T)
359 + ic2 = ic2 + w2_bias[i].float()
360 + out[mask] = ic2
361 +
362 + return (out.view(B, topk, -1) * topk_weight.unsqueeze(-1)).sum(dim=1).to(a.dtype)
363 +
364 +
365 +def _prepack_mxfp4_experts(
366 + w: torch.Tensor, w_scale: torch.Tensor
367 +) -> tuple[torch.Tensor, torch.Tensor]:
368 + """VNNI-prepack MXFP4 weights and repack scales."""
369 + packed_w = torch.ops._C.convert_weight_packed(w)
370 + packed_s = torch.ops._C.convert_scale_packed(w_scale)
371 + return packed_w, packed_s
372 +
373 +
374 +MXFP4_NUM_TOKENS = [1, 2, 32, 121]
375 +MXFP4_MOE_CONFIGS = [
376 + (128, 128, 4, 2),
377 + (256, 256, 8, 4),
378 + (352, 256, 8, 4),
379 + (512, 320, 8, 4),
380 +]
381 +
382 +
383 +@pytest.mark.parametrize("M", MXFP4_NUM_TOKENS)
384 +@pytest.mark.parametrize("N,K,E,topk", MXFP4_MOE_CONFIGS)
385 +@pytest.mark.parametrize("seed", [0])
386 +def test_mxfp4_cpu_fused_moe(M, N, K, E, topk, seed):
387 + """Test fused_experts_mxfp4_cpu against dequantized torch reference."""
388 + set_random_seed(seed)
389 + dtype = torch.bfloat16
390 +
391 + a = torch.randn(M, K, dtype=dtype) / 10
392 +
393 + # Generate and quantize weights
394 + w1_bf16 = torch.randn(E, 2 * N, K, dtype=dtype) / 10
395 + w1q, w1s = MXFP4QuantizeUtil.quantize(w1_bf16)
396 + w1s = w1s.reshape(E, 2 * N, K // 32)
397 + w1dq = MXFP4QuantizeUtil.dequantize(w1q, dtype, w1s)
398 +
399 + w2_bf16 = torch.randn(E, K, N, dtype=dtype) / 10
400 + w2q, w2s = MXFP4QuantizeUtil.quantize(w2_bf16)
401 + w2s = w2s.reshape(E, K, N // 32)
402 + w2dq = MXFP4QuantizeUtil.dequantize(w2q, dtype, w2s)
403 +
404 + # Routing
405 + score = torch.randn(M, E, dtype=dtype)
406 + score = torch.softmax(score, dim=-1, dtype=torch.float32)
407 + topk_weight, topk_ids = torch.topk(score, topk)
408 + topk_ids = topk_ids.to(torch.int32)
409 +
410 + # Reference
411 + ref_out = ref_mxfp4_fused_moe(a, w1dq, w2dq, topk_weight, topk_ids, topk)
412 +
413 + # Pack weights for kernel
414 + pw1, pw1s = _prepack_mxfp4_experts(w1q, w1s)
415 + pw2, pw2s = _prepack_mxfp4_experts(w2q, w2s)
416 +
417 + # Kernel
418 + out = ops.fused_experts_cpu(
419 + a.clone(),
420 + pw1,
421 + pw2,
422 + topk_weight,
423 + topk_ids,
424 + False, # inplace
425 + ops.CPUQuantMethod.MXFP4,
426 + pw1s, # w1_scale
427 + pw2s, # w2_scale
428 + None, # w1_zero
429 + None, # w2_zero
430 + None, # block_size
431 + )
432 +
433 + torch.testing.assert_close(ref_out.bfloat16(), out, atol=1e-2, rtol=1e-2)
434 +
435 +
436 +@pytest.mark.parametrize("M", [1, 32])
437 +@pytest.mark.parametrize("N,K,E,topk", [(128, 128, 4, 2), (64, 64, 4, 2)])
438 +@pytest.mark.parametrize("seed", [0])
439 +def test_mxfp4_cpu_fused_moe_bias_swiglu(M, N, K, E, topk, seed):
440 + """Test fused_experts_mxfp4_cpu with bias and SwigLU activation (GPT-OSS)."""
441 + set_random_seed(seed)
442 + dtype = torch.bfloat16
443 + alpha = 1.702
444 + limit = 7.0
445 +
446 + a = torch.randn(M, K, dtype=dtype) / 10
447 +
448 + # Generate and quantize weights
449 + w1_bf16 = torch.randn(E, 2 * N, K, dtype=dtype) / 10
450 + w1q, w1s = MXFP4QuantizeUtil.quantize(w1_bf16)
451 + w1s = w1s.reshape(E, 2 * N, K // 32)
452 + w1dq = MXFP4QuantizeUtil.dequantize(w1q, dtype, w1s)
453 + w1_b = torch.randn(E, 2 * N, dtype=torch.float32) / 10
454 +
455 + w2_bf16 = torch.randn(E, K, N, dtype=dtype) / 10
456 + w2q, w2s = MXFP4QuantizeUtil.quantize(w2_bf16)
457 + w2s = w2s.reshape(E, K, N // 32)
458 + w2dq = MXFP4QuantizeUtil.dequantize(w2q, dtype, w2s)
459 + w2_b = torch.randn(E, K, dtype=torch.float32) / 10
460 +
461 + # Routing
462 + score = torch.randn(M, E, dtype=dtype)
463 + score = torch.softmax(score, dim=-1, dtype=torch.float32)
464 + topk_weight, topk_ids = torch.topk(score, topk)
465 + topk_ids = topk_ids.to(torch.int32)
466 +
467 + # Reference
468 + ref_out = ref_mxfp4_fused_moe_gptoss(
469 + a, w1dq, w2dq, w1_b, w2_b, topk_weight, topk_ids, alpha, limit
470 + )
471 +
472 + # Pack weights for kernel
473 + pw1, pw1s = _prepack_mxfp4_experts(w1q, w1s)
474 + pw2, pw2s = _prepack_mxfp4_experts(w2q, w2s)
475 +
476 + # Kernel
477 + out = ops.fused_experts_cpu(
478 + a.clone(),
479 + pw1,
480 + pw2,
481 + topk_weight,
482 + topk_ids,
483 + False, # inplace
484 + ops.CPUQuantMethod.MXFP4,
485 + pw1s, # w1_scale
486 + pw2s, # w2_scale
487 + None, # w1_zero
488 + None, # w2_zero
489 + None, # block_size
490 + w1_bias=w1_b,
491 + w2_bias=w2_b,
492 + alpha=alpha,
493 + limit=limit,
494 + )
495 +
496 + torch.testing.assert_close(ref_out.bfloat16(), out, atol=1e-2, rtol=1e-2)
497 +
498 +
499 +if __name__ == "__main__":
500 + sys.exit(pytest.main([__file__, "-v"]))
tests/quantization/test_cpu_wna16.py
+1
@@ -15,6 +15,7 @@ MODELS = [
15 "OPEA/Qwen2.5-0.5B-Instruct-int4-sym-inc",
16 "Qwen/Qwen3-0.6B-FP8", # FP8 W8A16 block-quantized linear
17 "Qwen/Qwen3-30B-A3B-FP8", # FP8 W8A16 block-quantized MoE
18 + "openai/gpt-oss-20b", # MXFP4 W4A16
19 ]
20 DTYPE = ["bfloat16"]
21
vllm/_custom_ops.py
+14 -1
@@ -3195,6 +3195,7 @@ class CPUQuantMethod(IntEnum):
3195 INT8_W8A8 = 1
3196 FP8_W8A16 = 2
3197 INT4_W4A8 = 3
3198 + MXFP4 = 4
3199
3200
3201 if hasattr(torch.ops._C, "fused_experts_cpu"):
@@ -3213,6 +3214,10 @@ if hasattr(torch.ops._C, "fused_experts_cpu"):
3214 w1_zero: torch.Tensor | None,
3215 w2_zero: torch.Tensor | None,
3216 block_size: list[int] | None,
3217 + w1_bias: torch.Tensor | None,
3218 + w2_bias: torch.Tensor | None,
3219 + alpha: float | None,
3220 + limit: float | None,
3221 is_vnni: bool,
3222 ) -> torch.Tensor:
3223 return torch.empty_like(hidden_states)
@@ -3231,7 +3236,11 @@ def fused_experts_cpu(
3236 w1_zero: torch.Tensor | None,
3237 w2_zero: torch.Tensor | None,
3238 block_size: list[int] | None,
3234 - is_vnni: bool,
3239 + w1_bias: torch.Tensor | None = None,
3240 + w2_bias: torch.Tensor | None = None,
3241 + alpha: float | None = None,
3242 + limit: float | None = None,
3243 + is_vnni: bool = True,
3244 ) -> torch.Tensor:
3245 return torch.ops._C.fused_experts_cpu(
3246 hidden_states,
@@ -3246,6 +3255,10 @@ def fused_experts_cpu(
3255 w1_zero,
3256 w2_zero,
3257 block_size,
3258 + w1_bias,
3259 + w2_bias,
3260 + alpha,
3261 + limit,
3262 is_vnni,
3263 )
3264
vllm/model_executor/layers/fused_moe/cpu_fused_moe.py
+4
@@ -213,6 +213,10 @@ class SGLFusedMOE:
213 None, # w1_zero
214 None, # w2_zero
215 None, # block_size
216 + None, # w1_bias
217 + None, # w2_bias
218 + None, # alpha
219 + None, # limit
220 True, # is_vnni
221 )
222
vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
+160 -9
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""CPU FP8 W8A16 block-quantized fused MoE experts."""
3 +"""CPU FP8 W8A16 and MXFP4 W4A16 fused MoE experts."""
4
5 import torch
6
@@ -17,6 +17,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
17 QuantKey,
18 kFp8Dynamic128Sym,
19 kFp8Static128BlockSym,
20 + kMxfp4Static,
21 )
22 from vllm.platforms import current_platform
23
@@ -26,14 +27,8 @@ def prepare_fp8_moe_layer_for_cpu(
27 w2: torch.Tensor,
28 ) -> tuple[torch.Tensor, torch.Tensor]:
29 """VNNI-prepack FP8 MoE weights for CPU kernel."""
29 - num_experts = w13.size(0)
30 - packed_w13_list = []
31 - packed_w2_list = []
32 - for i in range(num_experts):
33 - packed_w13_list.append(torch.ops._C.convert_weight_packed(w13[i]))
34 - packed_w2_list.append(torch.ops._C.convert_weight_packed(w2[i]))
35 - packed_w13 = torch.stack(packed_w13_list)
36 - packed_w2 = torch.stack(packed_w2_list)
30 + packed_w13 = torch.ops._C.convert_weight_packed(w13)
31 + packed_w2 = torch.ops._C.convert_weight_packed(w2)
32 return packed_w13, packed_w2
33
34
@@ -171,5 +166,161 @@ class CPUExpertsFp8(mk.FusedMoEExpertsMonolithic):
166 None, # w1_zero
167 None, # w2_zero
168 block_shape, # block_size
169 + None, # w1_bias
170 + None, # w2_bias
171 + None, # alpha
172 + None, # limit
173 + True, # is_vnni
174 + )
175 +
176 +
177 +def prepare_mxfp4_moe_layer_for_cpu(
178 + w13: torch.Tensor,
179 + w2: torch.Tensor,
180 + w13_scale: torch.Tensor,
181 + w2_scale: torch.Tensor,
182 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
183 + """VNNI-prepack MXFP4 MoE weights and repack scales for CPU AMX kernel."""
184 + packed_w13 = torch.ops._C.convert_weight_packed(w13)
185 + packed_w2 = torch.ops._C.convert_weight_packed(w2)
186 + packed_w13_scale = torch.ops._C.convert_scale_packed(w13_scale)
187 + packed_w2_scale = torch.ops._C.convert_scale_packed(w2_scale)
188 + return packed_w13, packed_w2, packed_w13_scale, packed_w2_scale
189 +
190 +
191 +class CPUExpertsMxfp4(mk.FusedMoEExpertsMonolithic):
192 + """CPU MXFP4 W4A16 monolithic MoE experts."""
193 +
194 + def __init__(
195 + self,
196 + moe_config: FusedMoEConfig,
197 + quant_config: FusedMoEQuantConfig,
198 + ):
199 + super().__init__(
200 + moe_config,
201 + quant_config,
202 + )
203 +
204 + @property
205 + def expects_unquantized_inputs(self) -> bool:
206 + return True
207 +
208 + @staticmethod
209 + def activation_format() -> mk.FusedMoEActivationFormat:
210 + return mk.FusedMoEActivationFormat.Standard
211 +
212 + @staticmethod
213 + def _supports_current_device() -> bool:
214 + return current_platform.is_cpu()
215 +
216 + @staticmethod
217 + def _supports_no_act_and_mul() -> bool:
218 + return False
219 +
220 + @staticmethod
221 + def _supports_activation(activation: MoEActivation) -> bool:
222 + return activation in (MoEActivation.SILU, MoEActivation.SWIGLUOAI)
223 +
224 + @staticmethod
225 + def _supports_parallel_config(
226 + moe_parallel_config: FusedMoEParallelConfig,
227 + ) -> bool:
228 + return True
229 +
230 + @staticmethod
231 + def _supports_quant_scheme(
232 + weight_key: QuantKey | None,
233 + activation_key: QuantKey | None,
234 + ) -> bool:
235 + SUPPORTED_W_A = [
236 + (kMxfp4Static, None),
237 + ]
238 + return (weight_key, activation_key) in SUPPORTED_W_A
239 +
240 + @staticmethod
241 + def _supports_routing_method(
242 + routing_method: RoutingMethodType,
243 + weight_key: QuantKey | None,
244 + activation_key: QuantKey | None,
245 + ) -> bool:
246 + return routing_method in [
247 + RoutingMethodType.Default,
248 + RoutingMethodType.Renormalize,
249 + RoutingMethodType.RenormalizeNaive,
250 + ]
251 +
252 + @staticmethod
253 + def _supports_router_logits_dtype(
254 + router_logits_dtype: torch.dtype | None,
255 + routing_method: RoutingMethodType,
256 + ) -> bool:
257 + return True
258 +
259 + def supports_expert_map(self) -> bool:
260 + return False
261 +
262 + def apply(
263 + self,
264 + hidden_states: torch.Tensor,
265 + w1: torch.Tensor,
266 + w2: torch.Tensor,
267 + router_logits: torch.Tensor,
268 + activation: MoEActivation,
269 + global_num_experts: int,
270 + expert_map: torch.Tensor | None,
271 + a1q_scale: torch.Tensor | None,
272 + apply_router_weight_on_input: bool,
273 + # grouped topk + fused topk bias parameters
274 + num_expert_group: int | None = None,
275 + e_score_correction_bias: torch.Tensor | None = None,
276 + routed_scaling_factor: float | None = None,
277 + topk_group: int | None = None,
278 + ) -> torch.Tensor:
279 + from vllm.model_executor.layers.fused_moe.cpu_fused_moe import (
280 + select_experts,
281 + )
282 +
283 + topk_weights, topk_ids = select_experts(
284 + hidden_states=hidden_states,
285 + router_logits=router_logits,
286 + use_grouped_topk=num_expert_group is not None,
287 + top_k=self.moe_config.experts_per_token,
288 + renormalize=self.moe_config.routing_method
289 + in (
290 + RoutingMethodType.Renormalize,
291 + RoutingMethodType.RenormalizeNaive,
292 + ),
293 + topk_group=topk_group,
294 + num_expert_group=num_expert_group,
295 + scoring_func="softmax",
296 + routed_scaling_factor=(
297 + routed_scaling_factor if routed_scaling_factor is not None else 1.0
298 + ),
299 + e_score_correction_bias=e_score_correction_bias,
300 + )
301 +
302 + # Get bias and swiglu params from quant config
303 + w1_bias = self.quant_config.w1_bias
304 + w2_bias = self.quant_config.w2_bias
305 + alpha = getattr(self.quant_config, "gemm1_alpha", None)
306 + limit = getattr(self.quant_config, "gemm1_clamp_limit", None)
307 +
308 + return fused_experts_cpu(
309 + hidden_states,
310 + w1,
311 + w2,
312 + topk_weights,
313 + topk_ids,
314 + False, # inplace
315 + CPUQuantMethod.MXFP4, # moe_comp_method
316 + self.w1_scale, # w1_scale
317 + self.w2_scale, # w2_scale
318 + None, # w1_zero
319 + None, # w2_zero
320 + None, # block_size
321 + w1_bias,
322 + w2_bias,
323 + alpha,
324 + limit,
325 True, # is_vnni
326 )
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+49
@@ -79,6 +79,8 @@ class Mxfp4MoeBackend(Enum):
79 TRITON_UNFUSED = "TRITON_UNFUSED"
80 # XPU
81 XPU = "XPU"
82 + # CPU
83 + CPU = "CPU"
84 # Emulation
85 EMULATION = "EMULATION"
86 # Humming
@@ -198,6 +200,11 @@ def backend_to_kernel_cls(
200
201 return [XPUExpertsMXFp4]
202
203 + elif backend == Mxfp4MoeBackend.CPU:
204 + from vllm.model_executor.layers.fused_moe.experts.cpu_moe import CPUExpertsMxfp4
205 +
206 + return [CPUExpertsMxfp4]
207 +
208 elif backend == Mxfp4MoeBackend.EMULATION:
209 from vllm.model_executor.layers.fused_moe.experts.ocp_mx_emulation_moe import (
210 OCP_MXQuantizationEmulationTritonExperts,
@@ -237,6 +244,7 @@ def map_mxfp4_backend(runner_backend: MoEBackend) -> list[Mxfp4MoeBackend]:
244 ],
245 "aiter_mxfp4_fp8": [Mxfp4MoeBackend.AITER_MXFP4_FP8],
246 "xpu": [Mxfp4MoeBackend.XPU],
247 + "cpu": [Mxfp4MoeBackend.CPU],
248 "emulation": [Mxfp4MoeBackend.EMULATION],
249 }
250 if backends := mapping.get(runner_backend):
@@ -535,6 +543,17 @@ def select_mxfp4_moe_backend(
543 activation_format,
544 )
545
546 + if current_platform.is_cpu():
547 + backend = Mxfp4MoeBackend.CPU
548 + logger.info_once(_make_log_backend(backend))
549 + return _return_or_raise(
550 + Mxfp4MoeBackend.CPU,
551 + config,
552 + kMxfp4Static,
553 + None,
554 + activation_format,
555 + )
556 +
557 if current_platform.is_cuda() or current_platform.is_rocm():
558 raise NotImplementedError(
559 "No MXFP4 MoE backend supports the deployment configuration."
@@ -639,6 +658,10 @@ def mxfp4_round_up_hidden_size_and_intermediate_size(
658 elif current_platform.is_rocm():
659 intermediate_size = round_up(intermediate_size, 256)
660 hidden_size = round_up(hidden_size, 256)
661 + elif backend == Mxfp4MoeBackend.CPU:
662 + # CPU AMX kernel uses BLOCK_N=32, align to 32
663 + intermediate_size = round_up(intermediate_size, 32)
664 + hidden_size = round_up(hidden_size, 32)
665 else:
666 intermediate_size = round_up(intermediate_size, 64)
667 return hidden_size, intermediate_size
@@ -1098,6 +1121,31 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
1121 w13_bias,
1122 w2_bias,
1123 )
1124 + elif mxfp4_backend == Mxfp4MoeBackend.CPU:
1125 + from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
1126 + prepare_mxfp4_moe_layer_for_cpu,
1127 + )
1128 +
1129 + packed_w13, packed_w2, packed_w13_scale, packed_w2_scale = (
1130 + prepare_mxfp4_moe_layer_for_cpu(
1131 + w13_weight.data,
1132 + w2_weight.data,
1133 + w13_weight_scale.data,
1134 + w2_weight_scale.data,
1135 + )
1136 + )
1137 + if w13_bias is not None:
1138 + w13_bias = w13_bias.data.to(torch.float32)
1139 + if w2_bias is not None:
1140 + w2_bias = w2_bias.data.to(torch.float32)
1141 + return (
1142 + packed_w13,
1143 + packed_w2,
1144 + packed_w13_scale,
1145 + packed_w2_scale,
1146 + w13_bias,
1147 + w2_bias,
1148 + )
1149 elif mxfp4_backend == Mxfp4MoeBackend.EMULATION:
1150 # No additional transformation needed for emulation backend,
1151 # weights are dequantized on the fly in the experts class.
@@ -1501,6 +1549,7 @@ def make_mxfp4_moe_quant_config(
1549 Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16,
1550 Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,
1551 Mxfp4MoeBackend.AITER_MXFP4_BF16,
1552 + Mxfp4MoeBackend.CPU,
1553 ):
1554 return mxfp4_w4a16_moe_quant_config(
1555 w1_bias=w1_bias,