@tduncan / Nvidia-Optimization-Stack-VLLM / commits / b346479065

[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>

Yuwen Zhou committed Sep 18, 2026 at 17:59 UTC b34647906550f6e9dc65c55ecfb36b697d5e1f43
20 files changed +2924 -46
cmake/cpu_extension.cmake
+5
@@ -488,6 +488,7 @@ set(VLLM_EXT_SRC
488 "csrc/moe/dynamic_4bit_int_moe_cpu.cpp"
489 "csrc/cpu/cpu_fused_moe.cpp"
490 "csrc/cpu/cpu_attn.cpp"
491 + "csrc/cpu/cpu_isa.cpp"
492 "csrc/cpu/torch_bindings.cpp")
493
494 if (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64" AND VLLM_RVV_VLEN AND
@@ -546,11 +547,13 @@ if (ENABLE_X86_ISA)
547 "csrc/cpu/sgl-kernels/gemm.cpp"
548 "csrc/cpu/sgl-kernels/gemm_int8.cpp"
549 "csrc/cpu/sgl-kernels/gemm_fp8.cpp"
550 + "csrc/cpu/sgl-kernels/gemm_fp8_w8a8.cpp"
551 "csrc/cpu/sgl-kernels/gemm_int4.cpp"
552 "csrc/cpu/sgl-kernels/moe.cpp"
553 "csrc/cpu/sgl-kernels/moe_int8.cpp"
554 "csrc/cpu/sgl-kernels/moe_int4.cpp"
555 "csrc/cpu/sgl-kernels/moe_fp8.cpp"
556 + "csrc/cpu/sgl-kernels/moe_fp8_w8a8.cpp"
557 "csrc/cpu/sgl-kernels/bmm.cpp"
558 "csrc/cpu/sgl-kernels/decode.cpp"
559 "csrc/cpu/sgl-kernels/extend.cpp"
@@ -571,6 +574,7 @@ if (ENABLE_X86_ISA)
574 "csrc/cpu/utils.cpp"
575 "csrc/cpu/spec_decode_utils.cpp"
576 "csrc/cpu/cpu_attn.cpp"
577 + "csrc/cpu/cpu_isa.cpp"
578 "csrc/cpu/dnnl_kernels.cpp"
579 "csrc/cpu/mamba_cpu.cpp"
580 "csrc/cpu/torch_bindings.cpp"
@@ -587,6 +591,7 @@ if (ENABLE_X86_ISA)
591 "csrc/cpu/utils.cpp"
592 "csrc/cpu/spec_decode_utils.cpp"
593 "csrc/cpu/cpu_attn.cpp"
594 + "csrc/cpu/cpu_isa.cpp"
595 "csrc/cpu/mamba_cpu.cpp"
596 "csrc/cpu/dnnl_kernels.cpp"
597 "csrc/cpu/torch_bindings.cpp"
csrc/cpu/cpu_attn.cpp
+2 -11
@@ -1,16 +1,7 @@
1 #include "cpu_attn_dispatch_generated.h"
2
3 -#if defined(__x86_64__) || defined(_M_X64)
4 - #include <cpuid.h>
5 -// Runtime check for AMX-FP8 (Diamond Rapids): CPUID leaf 7, subleaf 1, EAX[21].
6 -static bool runtime_has_amx_fp8() {
7 - unsigned int eax = 0, ebx = 0, ecx = 0, edx = 0;
8 - if (__get_cpuid_count(7, 1, &eax, &ebx, &ecx, &edx)) {
9 - return (eax >> 21) & 1u;
10 - }
11 - return false;
12 -}
13 -#endif // x86_64
3 +// Runtime check for AMX-FP8, implemented in cpu_isa.cpp.
4 +extern bool runtime_has_amx_fp8();
5
6 // Maps kv_cache_dtype string to Fp8KVCacheDataType enum.
7 // "auto" -> kAuto(0); "fp8"/"fp8_e4m3" -> kFp8E4M3; "fp8_e5m2" -> kFp8E5M2.
csrc/cpu/cpu_isa.cpp new
+20
@@ -0,0 +1,20 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +// AMX-FP8 is reported at CPUID leaf 0x1E, subleaf 1, EAX[4].
5 +#if defined(__x86_64__) || defined(_M_X64)
6 + #include <cpuid.h>
7 +
8 +bool runtime_has_amx_fp8() {
9 + unsigned int eax = 0, ebx = 0, ecx = 0, edx = 0;
10 + if (__get_cpuid_max(0, nullptr) < 0x1e) return false;
11 + if (__get_cpuid_count(0x1e, 1, &eax, &ebx, &ecx, &edx)) {
12 + return (eax >> 4) & 1u;
13 + }
14 + return false;
15 +}
16 +
17 +bool cpu_has_amx_fp8() { return runtime_has_amx_fp8(); }
18 +#else
19 +bool cpu_has_amx_fp8() { return false; }
20 +#endif
csrc/cpu/sgl-kernels/common.h
+83
@@ -12,11 +12,45 @@
12 #include <ATen/ATen.h>
13 #include <ATen/Dispatch.h>
14 #include <ATen/Parallel.h>
15 +#include <cstring>
16
17 #if defined(_OPENMP)
18 #include <omp.h>
19 #endif
20
21 +// Adapted from sglang: dispatch utilities for FP8 W8A8 kernels
22 +#define AT_DISPATCH_OUT_TYPES(TYPE, NAME, ...) \
23 + AT_DISPATCH_SWITCH( \
24 + TYPE, \
25 + NAME, \
26 + AT_PRIVATE_CASE_TYPE_USING_HINT(at::ScalarType::Float, out_t, __VA_ARGS__) \
27 + AT_PRIVATE_CASE_TYPE_USING_HINT(at::ScalarType::BFloat16, out_t, __VA_ARGS__) \
28 + AT_PRIVATE_CASE_TYPE_USING_HINT(at::ScalarType::Half, out_t, __VA_ARGS__))
29 +
30 +#define AT_DISPATCH_CASE_ENUM_NO_RETURN(VALUE, TYPE, HINT, ...) \
31 + case VALUE: { \
32 + constexpr TYPE HINT = VALUE; \
33 + __VA_ARGS__; \
34 + break; \
35 + }
36 +
37 +#define AT_DISPATCH_BOOL_NO_RETURN(VALUE, NAME, HINT, ...) \
38 + [&]() { \
39 + switch (VALUE) { \
40 + AT_DISPATCH_CASE_ENUM_NO_RETURN(true, bool, HINT, __VA_ARGS__) \
41 + AT_DISPATCH_CASE_ENUM_NO_RETURN(false, bool, HINT, __VA_ARGS__) \
42 + } \
43 + }()
44 +
45 +#define AT_DISPATCH_QUANT_MODE_NO_RETURN(MODE, NAME, HINT, ...) \
46 + [&]() { \
47 + switch (MODE) { \
48 + AT_DISPATCH_CASE_ENUM_NO_RETURN(PER_TENSOR, int, HINT, __VA_ARGS__) \
49 + AT_DISPATCH_CASE_ENUM_NO_RETURN(PER_ROW, int, HINT, __VA_ARGS__) \
50 + AT_DISPATCH_CASE_ENUM_NO_RETURN(PER_GROUP, int, HINT, __VA_ARGS__) \
51 + } \
52 + }()
53 +
54 namespace {
55
56 // dispatch bool
@@ -472,4 +506,53 @@ inline T* conditional_data_ptr(const std::optional<at::Tensor>& opt) {
506 return opt.has_value() ? opt.value().data_ptr<T>() : nullptr;
507 }
508
509 +// Adapted from sglang: M blocking utilities for FP8 W8A8 kernels
510 +inline int64_t get_m_block(int64_t M) {
511 + if (M <= 48) {
512 + return M;
513 + } else if (M < 64) {
514 + return 32;
515 + } else if (M < 96) {
516 + return 64;
517 + } else {
518 + return 128;
519 + }
520 +}
521 +
522 +inline std::tuple<bool, int64_t, int64_t, int64_t> get_m_blocking(int64_t M) {
523 + bool parallel_on_M = M > 128;
524 + int64_t block_m = get_m_block(M);
525 + int64_t Mc = (M + block_m - 1) / block_m;
526 + int64_t Mc_parallel = parallel_on_M ? Mc : 1;
527 + return std::make_tuple(parallel_on_M, block_m, Mc, Mc_parallel);
528 +}
529 +
530 +#if defined(CPU_CAPABILITY_AVX512)
531 +template <typename T>
532 +void zero_buffer(T* data, int64_t size) {
533 + const int32_t vec_size = at::vec::Vectorized<T>::size();
534 + auto zero_vec = at::vec::Vectorized<T>(0);
535 + int64_t d = 0;
536 + for (; d < size - (size % vec_size); d += vec_size) {
537 + zero_vec.store(data + d);
538 + }
539 + if (d < size) {
540 + zero_vec.store(data + d, size - d);
541 + }
542 +}
543 +#else
544 +template <typename T>
545 +void zero_buffer(T* data, int64_t size) {
546 + memset(data, 0, sizeof(T) * size);
547 +}
548 +#endif
549 +
550 +#if defined(__x86_64__)
551 +bool avx10_2_available() {
552 + // __builtin_cpu_supports returns the masked feature bit rather than 0/1.
553 + static const bool available = __builtin_cpu_supports("avx10.2");
554 + return available;
555 +}
556 +#endif
557 +
558 } // anonymous namespace
csrc/cpu/sgl-kernels/gemm.h
+53 -1
@@ -94,7 +94,7 @@ enum class CPUActMethod : int {
94 gelu_and_mul = 2,
95 };
96
97 -enum class CPUQuantMethod : int64_t { BF16 = 0, INT8_W8A8 = 1, FP8_W8A16 = 2, INT4_W4A8 = 3, MXFP4 = 4 };
97 +enum class CPUQuantMethod : int64_t { BF16 = 0, INT8_W8A8 = 1, FP8_W8A16 = 2, INT4_W4A8 = 3, MXFP4 = 4, FP8_W8A8 = 5 };
98
99 constexpr bool operator==(CPUQuantMethod a, int64_t b) {
100 return static_cast<int64_t>(a) == b;
@@ -395,3 +395,55 @@ void tinygemm_kernel(
395 bool brg,
396 int64_t block_size_K,
397 bool do_unpack = true);
398 +
399 +// FP8 W8A8 tinygemm kernel (defined in gemm_fp8_w8a8.cpp)
400 +void tinygemm_kernel(
401 + float* C,
402 + const at::Float8_e4m3fn* A,
403 + const float* scales_a,
404 + const at::Float8_e4m3fn* B,
405 + const float* scales_b,
406 + int64_t M,
407 + int64_t K,
408 + int64_t lda,
409 + int64_t ldc,
410 + int64_t ldsa,
411 + float* ukernel_buf,
412 + at::BFloat16* dqA_buf,
413 + at::BFloat16* dqB_buf);
414 +
415 +// FP8 W8A8 quantize activation (defined in gemm_fp8_w8a8.cpp)
416 +std::tuple<at::Tensor, at::Tensor> quantize_fp8e4m3_vec(
417 + const at::Tensor& t,
418 + bool channelwise,
419 + c10::optional<at::Tensor> scale_opt);
420 +
421 +// FP8 W8A8 fused MoE kernel (defined in moe_fp8_w8a8.cpp)
422 +template <typename scalar_t>
423 +void fused_experts_fp8_a8_kernel_impl(
424 + scalar_t* __restrict__ output,
425 + scalar_t* __restrict__ ic0,
426 + scalar_t* __restrict__ ic1,
427 + scalar_t* __restrict__ ic2,
428 + at::Float8_e4m3fn* __restrict__ A_tmp,
429 + scalar_t* __restrict__ B_tmp,
430 + float* __restrict__ C_tmp,
431 + float* __restrict__ Ukernel_tmp,
432 + const at::Float8_e4m3fn* __restrict__ input,
433 + const at::Float8_e4m3fn* __restrict__ packed_w1,
434 + const at::Float8_e4m3fn* __restrict__ packed_w2,
435 + const float* __restrict__ As,
436 + const float* __restrict__ w1s,
437 + const float* __restrict__ w2s,
438 + int64_t block_size_N,
439 + int64_t block_size_K,
440 + const float* __restrict__ topk_weights,
441 + const int32_t* __restrict__ sorted_ids,
442 + const int32_t* __restrict__ expert_ids,
443 + const int32_t* __restrict__ offsets,
444 + int64_t M,
445 + int64_t N,
446 + int64_t K,
447 + int64_t E,
448 + int64_t topk,
449 + int64_t num_tokens_post_pad);
csrc/cpu/sgl-kernels/gemm_fp8_w8a8.cpp new
+1093
@@ -0,0 +1,1093 @@
1 +// Adapted from
2 +// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
3 +
4 +// clang-format off
5 +
6 +#include <ATen/native/CPUBlas.h>
7 +#include <c10/util/Unroll.h>
8 +#include <torch/all.h>
9 +
10 +#include "common.h"
11 +#include "vec.h"
12 +
13 +// Runtime check for AMX-FP8, implemented in cpu_isa.cpp.
14 +extern bool cpu_has_amx_fp8();
15 +
16 +namespace {
17 +
18 +#define BLOCK_N 32
19 +
20 +#define PER_TENSOR 1
21 +#define PER_ROW 2
22 +#define PER_GROUP 3
23 +
24 +static bool cpublas_checked = false;
25 +static bool cpublas_can_pack = false;
26 +
27 +bool cpublas_could_pack() {
28 + // the could_pack check requires AMX support implicitly
29 + if (cpublas_checked) {
30 + return cpublas_can_pack;
31 + }
32 + cpublas_can_pack = cpu_has_amx_fp8() && at::native::cpublas::could_pack(at::kFloat8_e4m3fn);
33 + cpublas_checked = true;
34 + return cpublas_can_pack;
35 +}
36 +
37 +#if defined(CPU_CAPABILITY_AVX512)
38 +static void cvt_f8e4m3_to_bf16(
39 + const at::Float8_e4m3fn* __restrict__ in, at::BFloat16* out, int64_t rows, int64_t cols, int64_t stride) {
40 + if (stride == cols) {
41 + // A contiguous buffer
42 + size_t len = rows * cols;
43 + size_t i = 0;
44 + for (; i < len; i += 32) {
45 + __m256i fp8_vec = _mm256_loadu_si256((__m256i*)&in[i]);
46 + __m512bh bf16_vec = cvt_e4m3_bf16_intrinsic_no_nan(fp8_vec);
47 + _mm512_storeu_si512((__m512i*)(out + i), (__m512i)bf16_vec);
48 + }
49 + for (; i < len; ++i) {
50 + out[i] = (at::BFloat16)in[i];
51 + }
52 + } else {
53 + // Non-contiguous. Access each row with stride
54 + TORCH_CHECK(stride > cols);
55 + for (int r = 0; r < rows; ++r) {
56 + size_t i = 0;
57 + size_t vec_len = cols / 32 * 32;
58 + for (; i < vec_len; i += 32) {
59 + __m256i fp8_vec = _mm256_loadu_si256((__m256i*)&in[r * stride + i]);
60 + __m512bh bf16_vec = cvt_e4m3_bf16_intrinsic_no_nan(fp8_vec);
61 + _mm512_storeu_si512((__m512i*)(out + r * cols + i), (__m512i)bf16_vec);
62 + }
63 + for (; i < cols; ++i) {
64 + out[r * cols + i] = (at::BFloat16)in[r * stride + i];
65 + }
66 + }
67 + }
68 +}
69 +
70 +// accumulate and store result to buffer
71 +// if act/wei are per_group quantized, apply scales
72 +template <bool accum, int64_t N, int act_quant_mode, int wei_quant_mode>
73 +static void _accumulate_result(
74 + float* __restrict__ output,
75 + const float* __restrict__ input,
76 + const float* __restrict__ scale_a,
77 + const float* __restrict__ scale_b,
78 + int M,
79 + int ldi,
80 + int ldo,
81 + int ldsa = 1) {
82 + float a_scale, b_scale;
83 + __m512 va_scale;
84 + __m512 vb_scale;
85 + for (int m = 0; m < M; ++m) {
86 + if constexpr (act_quant_mode == PER_GROUP) {
87 + a_scale = *(scale_a + m * ldsa);
88 + va_scale = _mm512_set1_ps(a_scale);
89 + }
90 + constexpr int N_UNROLL = N / 16;
91 + c10::ForcedUnroll<N_UNROLL>{}([&](auto i) {
92 + constexpr int n = i * 16;
93 + __m512 vc_f = _mm512_loadu_ps(input + m * ldi + n);
94 + if constexpr (act_quant_mode == PER_GROUP) {
95 + vc_f = _mm512_mul_ps(vc_f, va_scale);
96 + }
97 + if constexpr (wei_quant_mode == PER_GROUP) {
98 + vb_scale = _mm512_loadu_ps(scale_b + n);
99 + vc_f = _mm512_mul_ps(vc_f, vb_scale);
100 + }
101 + if constexpr (accum) {
102 + __m512 vo = _mm512_loadu_ps(output + m * ldo + n);
103 + _mm512_storeu_ps(output + m * ldo + n, _mm512_add_ps(vo, vc_f));
104 + } else {
105 + _mm512_storeu_ps(output + m * ldo + n, vc_f);
106 + }
107 + });
108 + constexpr int tail_start = N / 16 * 16;
109 + for (int n = tail_start; n < N; ++n) {
110 + float dq_val = input[m * ldi + n];
111 + if constexpr (act_quant_mode == PER_GROUP) {
112 + dq_val = dq_val * a_scale;
113 + }
114 + if constexpr (wei_quant_mode == PER_GROUP) {
115 + b_scale = scale_b[n];
116 + dq_val = dq_val * b_scale;
117 + }
118 + if constexpr (accum) {
119 + output[m * ldo + n] += dq_val;
120 + } else {
121 + output[m * ldo + n] = dq_val;
122 + }
123 + }
124 + }
125 +}
126 +
127 +// Store result to output buffer with dtype conversion
128 +// If act/wei are per_row or per_tensor quantized, apply scales
129 +// If bias is not null, add bias
130 +template <typename out_dtype, int64_t N, int act_quant_mode, int wei_quant_mode>
131 +inline void store_out(
132 + const float* y_buf,
133 + out_dtype* c_ptr,
134 + int64_t M,
135 + int64_t lda,
136 + const float* scales_a,
137 + const float* scales_b,
138 + const float* bias) {
139 + float a_scale = 1.0, b_scale = 1.0;
140 + __m512 va_scale, vb_scale;
141 + if constexpr (act_quant_mode == PER_TENSOR) {
142 + a_scale = *scales_a;
143 + }
144 + if constexpr (wei_quant_mode == PER_TENSOR) {
145 + b_scale = *scales_b;
146 + vb_scale = _mm512_set1_ps(b_scale);
147 + }
148 + for (int i = 0; i < M; ++i) {
149 + if constexpr (act_quant_mode == PER_ROW) {
150 + a_scale = *(scales_a + i);
151 + }
152 + if constexpr (act_quant_mode != PER_GROUP) {
153 + va_scale = _mm512_set1_ps(a_scale);
154 + }
155 + constexpr int N_UNROLL = N / 16;
156 + c10::ForcedUnroll<N_UNROLL>{}([&](auto idx) {
157 + constexpr int j = idx * 16;
158 + __m512 y_vec = _mm512_loadu_ps(y_buf + i * N + j);
159 + __m512 bias_vec = bias ? _mm512_loadu_ps(bias + j) : _mm512_setzero_ps();
160 + if constexpr (act_quant_mode != PER_GROUP) {
161 + y_vec = _mm512_mul_ps(y_vec, va_scale);
162 + }
163 + if constexpr (wei_quant_mode == PER_ROW) {
164 + vb_scale = _mm512_loadu_ps(scales_b + j);
165 + }
166 + if constexpr (wei_quant_mode != PER_GROUP) {
167 + y_vec = _mm512_mul_ps(y_vec, vb_scale);
168 + }
169 + y_vec = _mm512_add_ps(y_vec, bias_vec);
170 + if constexpr (std::is_same<out_dtype, float>::value) {
171 + _mm512_storeu_ps(c_ptr + i * lda + j, y_vec);
172 + } else if constexpr (std::is_same<out_dtype, at::BFloat16>::value) {
173 + __m256i y_bf16_vec = at::vec::cvtfp32_bf16(y_vec);
174 + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c_ptr + i * lda + j), y_bf16_vec);
175 + } else if constexpr (std::is_same<out_dtype, at::Half>::value) {
176 + __m256i y_fp16_vec = at::vec::cvtfp32_fp16(y_vec);
177 + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c_ptr + i * lda + j), y_fp16_vec);
178 + } else {
179 + TORCH_CHECK(false, "Unsupported output dtype");
180 + }
181 + });
182 + constexpr int tail_start = N / 16 * 16;
183 + for (int j = tail_start; j < N; ++j) {
184 + if constexpr (wei_quant_mode == PER_ROW) {
185 + b_scale = scales_b[j];
186 + }
187 + c_ptr[i * lda + j] = static_cast<out_dtype>(y_buf[i * N + j] * a_scale * b_scale);
188 + }
189 + } // for M
190 +}
191 +
192 +#else // no AVX512
193 +
194 +static void cvt_f8e4m3_to_bf16(
195 + const at::Float8_e4m3fn* __restrict__ in, at::BFloat16* out, int64_t rows, int64_t cols, int64_t stride) {
196 + for (int r = 0; r < rows; ++r) {
197 + for (int c = 0; c < cols; ++c) {
198 + out[r * cols + c] = (at::BFloat16)in[r * stride + c];
199 + }
200 + }
201 +}
202 +
203 +// Store result to output buffer with dtype conversion
204 +// If act/wei are per_row or per_tensor quantized, apply scales
205 +// If bias is not null, add bias
206 +template <typename out_dtype, int64_t N, int act_quant_mode, int wei_quant_mode>
207 +inline void store_out(
208 + const float* y_buf,
209 + out_dtype* c_ptr,
210 + int64_t M,
211 + int64_t lda,
212 + const float* scales_a,
213 + const float* scales_b,
214 + const float* bias) {
215 + float a_scale = 1.0, b_scale = 1.0;
216 + if constexpr (act_quant_mode == PER_TENSOR) {
217 + a_scale = *scales_a;
218 + }
219 + if constexpr (wei_quant_mode == PER_TENSOR) {
220 + b_scale = *scales_b;
221 + }
222 + for (int i = 0; i < M; ++i) {
223 + if constexpr (act_quant_mode == PER_ROW) {
224 + a_scale = *(scales_a + i);
225 + }
226 + for (int j = 0; j < N; ++j) {
227 + if constexpr (wei_quant_mode == PER_ROW) {
228 + b_scale = scales_b[j];
229 + }
230 + c_ptr[i * lda + j] = static_cast<out_dtype>(y_buf[i * N + j] * a_scale * b_scale);
231 + }
232 + } // for M
233 +}
234 +
235 +#endif // CPU_CAPABILITY_AVX512
236 +
237 +template <bool cpublas_can_pack, int64_t N, int act_quant_mode, int wei_quant_mode>
238 +void _micro_gemm(
239 + float* C,
240 + const at::Float8_e4m3fn* A,
241 + const float* scales_a,
242 + const at::Float8_e4m3fn* B,
243 + const float* scales_b,
244 + int64_t M,
245 + int64_t K,
246 + int64_t lda,
247 + int64_t ldc,
248 + int64_t ldsa,
249 + float* ukernel_buf,
250 + at::BFloat16* dqA_buf,
251 + at::BFloat16* dqB_buf) {
252 + // If FP8 brgemm is not available, convert A/B to bf16 for computation
253 + // Compute GEMM fp8 * fp8 -> fp32 (or bf16 * bf16 -> fp32)
254 + // If per_group quant, apply scales. Otherwise, don't apply scales here
255 + // Finally accumulate and store results
256 +#if defined(CPU_CAPABILITY_AVX512)
257 + if constexpr (cpublas_can_pack) {
258 +#ifdef CPUBLAS_BRGEMM_F8F8F32
259 + at::native::cpublas::brgemm(
260 + M, N, K, lda /*lda*/, N /*ldb*/, N /*ldc*/, false /* add_C */, A, B, ukernel_buf, true /* is_vnni */);
261 +#else
262 + cvt_f8e4m3_to_bf16(A, dqA_buf, M, K, lda);
263 + cvt_f8e4m3_to_bf16(B, dqB_buf, K, N, N);
264 + at::native::cpublas::brgemm(
265 + M, N, K, K /*lda*/, N /*ldb*/, N /*ldc*/, false /* add_C */, dqA_buf, dqB_buf, ukernel_buf, true /* is_vnni */);
266 +#endif
267 + _mm_prefetch(B + N * (K + 128), _MM_HINT_T0);
268 + _mm_prefetch(A + K + 128, _MM_HINT_T0);
269 + _accumulate_result<true, N, act_quant_mode, wei_quant_mode>(
270 + C, ukernel_buf, scales_a, scales_b, M, N /*ldi*/, ldc, ldsa);
271 + } else
272 +#endif
273 + {
274 + for (int64_t i = 0; i < M; ++i) {
275 + for (int64_t j = 0; j < N; ++j) {
276 + float sum = 0;
277 + for (int64_t k = 0; k < K; ++k) {
278 + sum += ((float)A[i * lda + k] * (float)B[k * N + j]);
279 + }
280 + if constexpr (act_quant_mode == PER_GROUP) {
281 + sum *= scales_a[i * ldsa];
282 + }
283 + if constexpr (wei_quant_mode == PER_GROUP) {
284 + sum *= scales_b[j];
285 + }
286 + C[i * ldc + j] += sum;
287 + }
288 + }
289 + }
290 +}
291 +
292 +template <typename out_dtype, bool cpublas_can_pack, int act_quant_mode, int wei_quant_mode>
293 +void _float8_linear_impl(
294 + const at::Tensor& input,
295 + const at::Tensor& input_scales,
296 + const at::Tensor& weight,
297 + const at::Tensor& weight_scales,
298 + const std::optional<at::Tensor>& bias,
299 + at::Tensor& output) {
300 + // input shape = [..., K]
301 + // input is per token quantized
302 + int64_t K = input.size(-1);
303 + auto input_view = input.view({-1, K});
304 + int64_t M = input_view.size(0);
305 +
306 + // weight shape = [Nc, Kc, block_k, block_n]
307 + // scales shape = [Nc, G, block_n]
308 + int64_t Nc = weight.size(0);
309 + int64_t Kc = wei_quant_mode != PER_GROUP ? 1 : weight.size(1);
310 + int64_t block_k = wei_quant_mode != PER_GROUP ? weight.size(1) * weight.size(2) : weight.size(2);
311 + constexpr int64_t block_n = BLOCK_N;
312 + TORCH_CHECK(weight.size(3) == block_n, "Float8 linear: unexpected weight shape");
313 + int64_t N = Nc * block_n;
314 + TORCH_CHECK(K == Kc * block_k, "Float8 linear: weight and input shapes mismatch");
315 + auto [parallel_on_M, block_m, Mc, Mc_parallel] = get_m_blocking(M);
316 +
317 + // scales shape = [Nc, G, block_n]
318 + int64_t num_groups = wei_quant_mode == PER_TENSOR ? 1 : weight_scales.size(1);
319 + TORCH_CHECK(K % num_groups == 0, "K should be divisible by num_groups");
320 + int64_t group_size = K / num_groups;
321 + TORCH_CHECK(group_size % block_k == 0, "Float8 linear: group_size should be divisible by block_k");
322 + int64_t block_per_group = group_size / block_k;
323 + TORCH_CHECK(
324 + input_scales.numel() == 1 || input_scales.numel() == M || input_scales.numel() == M * num_groups,
325 + "Float8 linear: unexpected input scales shape");
326 + auto ldsa = act_quant_mode == PER_TENSOR ? 0 : act_quant_mode == PER_ROW ? 1 : num_groups;
327 +
328 + const at::Float8_e4m3fn* a_ptr = input_view.data_ptr<at::Float8_e4m3fn>();
329 + const float* a_scales_ptr = input_scales.data_ptr<float>();
330 + const at::Float8_e4m3fn* b_ptr = weight.data_ptr<at::Float8_e4m3fn>();
331 + const float* b_scales_ptr = weight_scales.data_ptr<float>();
332 + out_dtype* c_ptr = output.data_ptr<out_dtype>();
333 + const float* bias_ptr = bias.has_value() ? bias.value().data_ptr<float>() : nullptr;
334 +
335 + int64_t block_size = block_m * block_n;
336 + int64_t num_thread = at::get_num_threads();
337 + at::Tensor y_buffer = at::empty({num_thread, block_size}, output.options().dtype(at::kFloat));
338 + // Create buffer for brgemm output and dqA/dqB (optional)
339 +#if defined(CPU_CAPABILITY_AVX512)
340 + // buffer for brgemm output in float32
341 + int64_t buffer_size = block_size * 2; // float32 = bfloat16 * 2
342 +#ifndef CPUBLAS_BRGEMM_F8F8F32
343 + // buffers for dqA & dqB in bf16
344 + buffer_size += (block_k * block_n + block_m * block_k);
345 +#endif
346 + at::Tensor micro_gemm_buffer = at::empty({num_thread, buffer_size}, output.options().dtype(at::kBFloat16));
347 +#endif
348 +
349 + parallel_2d(Mc, Nc, [&](int64_t mc0, int64_t mc1, int64_t nc0, int64_t nc1) {
350 + int tid = get_thread_num();
351 + float* y_buf = y_buffer.data_ptr<float>() + tid * block_size;
352 + at::BFloat16 *dqA_buffer = nullptr, *dqB_buffer = nullptr;
353 + float* ukernel_buf = nullptr;
354 +#if defined(CPU_CAPABILITY_AVX512)
355 + at::BFloat16* micro_gemm_buf = micro_gemm_buffer.data_ptr<at::BFloat16>() + tid * buffer_size;
356 + ukernel_buf = reinterpret_cast<float*>(micro_gemm_buf);
357 +#ifndef CPUBLAS_BRGEMM_F8F8F32
358 + dqA_buffer = micro_gemm_buf;
359 + dqB_buffer = micro_gemm_buf + block_m * block_k;
360 + ukernel_buf = reinterpret_cast<float*>(micro_gemm_buf + block_m * block_k + block_k * block_n);
361 +#endif
362 +#endif
363 +
364 + loop_2d<at::Float8_e4m3fn>(mc0, mc1, nc0, nc1, block_n * K, [&](int64_t mci, int64_t nc, int64_t) {
365 + int64_t m_size = mci * block_m + block_m > M ? M - mci * block_m : block_m;
366 + zero_buffer(y_buf, m_size * block_n);
367 + for (int kci = 0; kci < Kc; ++kci) {
368 + auto scales_a = a_scales_ptr + mci * block_m * num_groups + kci / block_per_group;
369 + auto scales_b = b_scales_ptr + nc * block_n * num_groups + kci / block_per_group * block_n;
370 + _micro_gemm<cpublas_can_pack, block_n, act_quant_mode, wei_quant_mode>(
371 + /* C */ y_buf,
372 + /* A */ a_ptr + mci * block_m * K + kci * block_k,
373 + /* scales_a */ scales_a,
374 + /* B */ b_ptr + (nc * Kc + kci) * block_n * block_k,
375 + /* scales_b */ scales_b,
376 + /* M */ m_size,
377 + /* K */ block_k,
378 + /* lda */ K,
379 + /* ldc */ block_n,
380 + /* ldsa */ ldsa,
381 + /* ukernel_buf */ ukernel_buf,
382 + /* dqA_buf */ dqA_buffer,
383 + /* dqB_buf */ dqB_buffer);
384 + }
385 + auto scales_a = act_quant_mode == PER_TENSOR ? a_scales_ptr
386 + : act_quant_mode == PER_ROW ? a_scales_ptr + mci * block_m
387 + : nullptr;
388 + auto scales_b = wei_quant_mode == PER_TENSOR ? b_scales_ptr
389 + : wei_quant_mode == PER_ROW ? b_scales_ptr + nc * block_n
390 + : nullptr;
391 + auto bias_data = bias_ptr ? bias_ptr + nc * block_n : nullptr;
392 + store_out<out_dtype, block_n, act_quant_mode, wei_quant_mode>(
393 + y_buf, c_ptr + mci * block_m * N + nc * block_n, m_size, N /*lda*/, scales_a, scales_b, bias_data);
394 + });
395 +
396 + if constexpr (cpublas_can_pack) {
397 + at::native::cpublas::brgemm_release();
398 + }
399 + });
400 +}
401 +
402 +} // anonymous namespace
403 +
404 +void tinygemm_kernel(
405 + float* C,
406 + const at::Float8_e4m3fn* A,
407 + const float* scales_a,
408 + const at::Float8_e4m3fn* B,
409 + const float* scales_b,
410 + int64_t M,
411 + int64_t K,
412 + int64_t lda,
413 + int64_t ldc,
414 + int64_t ldsa,
415 + float* ukernel_buf,
416 + at::BFloat16* dqA_buf,
417 + at::BFloat16* dqB_buf) {
418 + // act_quant_mode = per row (2), wei_quant_mode = per group (3)
419 + if (cpublas_could_pack()) {
420 + _micro_gemm<true, BLOCK_N, 2, 3>(C, A, scales_a, B, scales_b, M, K, lda, ldc, ldsa, ukernel_buf, dqA_buf, dqB_buf);
421 + } else {
422 + _micro_gemm<false, BLOCK_N, 2, 3>(C, A, scales_a, B, scales_b, M, K, lda, ldc, ldsa, ukernel_buf, dqA_buf, dqB_buf);
423 + }
424 +}
425 +
426 +#define INSTANTIATE_TINYGEMM_TEMPLATE() \
427 + void tinygemm_kernel( \
428 + float* C, \
429 + const at::Float8_e4m3fn* A, \
430 + const float* scales_a, \
431 + const at::Float8_e4m3fn* B, \
432 + const float* scales_b, \
433 + int64_t M, \
434 + int64_t K, \
435 + int64_t lda, \
436 + int64_t ldc, \
437 + int64_t ldsa, \
438 + float* ukernel_buf, \
439 + at::BFloat16* dqA_buf, \
440 + at::BFloat16* dqB_buf)
441 +
442 +INSTANTIATE_TINYGEMM_TEMPLATE();
443 +/*
444 +return: packed_weight, packed_scales
445 +*/
446 +std::tuple<at::Tensor, at::Tensor> float8_linear_prepack_impl(const at::Tensor& weight, const at::Tensor& scales) {
447 + // weight shape = [N, K]
448 + // scales shape = [N, G]
449 + TORCH_CHECK(weight.dim() == 2, "Float8 linear CPU: Weight should be a 2D tensor for packing");
450 + TORCH_CHECK(weight.size(1) % 2 == 0, "Float8 linear CPU: Weight should have even number of columns for packing");
451 +
452 + auto new_scales = scales;
453 + if (new_scales.dim() == 1) {
454 + new_scales.unsqueeze_(1);
455 + }
456 + new_scales = new_scales.to(at::kFloat);
457 + int N = weight.size(0);
458 + int K = weight.size(1);
459 + int G = scales.size(1);
460 + int group_size = K / G;
461 + int block_k = group_size > 128 ? 128 : group_size;
462 + while (K % block_k != 0) {
463 + block_k /= 2;
464 + }
465 + TORCH_CHECK(
466 + block_k > 0 && block_k <= group_size, "Float8 linear CPU: Invalid block_k size, should be in (0, group_size]");
467 + constexpr int block_n = BLOCK_N;
468 + int Nc = N / block_n;
469 + int Kc = K / block_k;
470 +
471 + // Reorder weight to [N/block_n, K/block_k, block_k, block_n]
472 + // Reorder scales to [N/block_n, G, block_n]
473 + auto weight_view = weight.view({Nc, block_n, Kc, block_k});
474 + at::Tensor weight_reordered = weight_view.permute({0, 2, 3, 1}).contiguous();
475 + at::Tensor blocked_weight;
476 + at::Tensor blocked_scales = new_scales.view({Nc, block_n, G}).permute({0, 2, 1}).contiguous();
477 +
478 +#if defined(CPU_CAPABILITY_AVX512)
479 + if (cpublas_could_pack()) {
480 +#ifdef CPUBLAS_BRGEMM_F8F8F32
481 + constexpr int vnni_size = 4; // for fp8
482 +#else
483 + constexpr int vnni_size = 2; // for float16
484 +#endif
485 + blocked_weight = at::empty({Nc, Kc, block_k, block_n}, weight.options());
486 + auto weight_ptr = reinterpret_cast<uint8_t*>(weight_reordered.data_ptr());
487 + auto blocked_weight_ptr = reinterpret_cast<uint8_t*>(blocked_weight.data_ptr());
488 + int64_t num_blocks = Nc * Kc;
489 + at::parallel_for(0, num_blocks, 1, [&](int64_t begin, int64_t end) {
490 + for (const auto i : c10::irange(begin, end)) {
491 + auto in_ptr = weight_ptr + i * block_k * block_n;
492 + auto out_ptr = blocked_weight_ptr + i * block_k * block_n;
493 +
494 + // Reorder weight block to VNNI
495 + // plain shape = [block_k, block_n]
496 + // packed shape = [block_k / VNNI_SIZE, block_n, VNNI_SIZE] viewed as [block_k, block_n]
497 + constexpr int n_group_size = 8;
498 + constexpr int n_group = block_n / n_group_size; // 4
499 + for (int nb = 0; nb < n_group; ++nb) {
500 + for (int k = 0; k < block_k; k += vnni_size) {
501 + for (int ni = 0; ni < n_group_size; ++ni) {
502 + for (int ki = 0; ki < vnni_size; ++ki) {
503 + int src_idx = nb * n_group_size + ni + (k + ki) * block_n;
504 + int dst_idx = (nb * n_group_size + ni) * vnni_size + k * block_n + ki;
505 + *(out_ptr + dst_idx) = *(in_ptr + src_idx);
506 + }
507 + }
508 + }
509 + }
510 + }
511 + });
512 + } else
513 +#endif
514 + {
515 + blocked_weight = weight_reordered;
516 + }
517 +
518 + return std::make_tuple(std::move(blocked_weight), std::move(blocked_scales));
519 +}
520 +
521 +// AVX512 optimized channel-wise max computation
522 +inline void
523 +compute_channel_max_avx512(const float* data, float* max_vals, int64_t num_channels, int64_t elements_per_channel) {
524 + at::parallel_for(0, num_channels, 1, [&](int64_t start, int64_t end) {
525 + for (int64_t c = start; c < end; ++c) {
526 + const float* channel_data = data + c * elements_per_channel;
527 + float max_val = 0.0f;
528 +
529 + int64_t i = 0;
530 + // Process 16 elements at a time
531 + __m512 max_vec = _mm512_setzero_ps();
532 +
533 + for (; i <= elements_per_channel - 16; i += 16) {
534 + __m512 data_vec = _mm512_loadu_ps(&channel_data[i]);
535 + __m512 abs_vec = _mm512_abs_ps(data_vec);
536 + max_vec = _mm512_max_ps(max_vec, abs_vec);
537 + }
538 +
539 + // Horizontal max of the vector
540 + float hmax = _mm512_reduce_max_ps(max_vec);
541 + max_val = std::max(max_val, hmax);
542 +
543 + // Handle remaining elements
544 + for (; i < elements_per_channel; ++i) {
545 + max_val = std::max(max_val, std::abs(channel_data[i]));
546 + }
547 +
548 + max_vals[c] = max_val;
549 + }
550 + });
551 +}
552 +
553 +// AVX512 optimized scaling and clamping for channel-wise quantization
554 +inline void scale_clamp_channelwise_avx512(
555 + const float* src,
556 + float* dst,
557 + const float* scales,
558 + int64_t num_channels,
559 + int64_t elements_per_channel,
560 + float quant_max,
561 + float neg_quant_max) {
562 + at::parallel_for(0, num_channels, 1, [&](int64_t start, int64_t end) {
563 + for (int64_t c = start; c < end; ++c) {
564 + float scale_val = scales[c];
565 + float scale_reciprocal = 1.0f / scale_val;
566 +
567 + const float* channel_src = src + c * elements_per_channel;
568 + float* channel_dst = dst + c * elements_per_channel;
569 +
570 + int64_t i = 0;
571 + const __m512 scale_recip_vec = _mm512_set1_ps(scale_reciprocal);
572 + const __m512 quant_max_vec = _mm512_set1_ps(quant_max);
573 + const __m512 neg_quant_max_vec = _mm512_set1_ps(neg_quant_max);
574 +
575 + for (; i <= elements_per_channel - 16; i += 16) {
576 + __m512 src_vec = _mm512_loadu_ps(&channel_src[i]);
577 + __m512 scaled_vec = _mm512_mul_ps(src_vec, scale_recip_vec);
578 + __m512 clamped_vec = _mm512_min_ps(_mm512_max_ps(scaled_vec, neg_quant_max_vec), quant_max_vec);
579 + _mm512_storeu_ps(&channel_dst[i], clamped_vec);
580 + }
581 +
582 + // Handle remaining elements
583 + for (; i < elements_per_channel; ++i) {
584 + float scaled = channel_src[i] * scale_reciprocal;
585 + channel_dst[i] = std::clamp(scaled, neg_quant_max, quant_max);
586 + }
587 + }
588 + });
589 +}
590 +
591 +// AVX512 optimized scaling and clamping for global quantization
592 +inline void scale_clamp_global_avx512(
593 + const float* src, float* dst, int64_t size, float scale_reciprocal, float quant_max, float neg_quant_max) {
594 + if (size <= 0) return;
595 +
596 + const __m512 scale_recip_vec = _mm512_set1_ps(scale_reciprocal);
597 + const __m512 quant_max_vec = _mm512_set1_ps(quant_max);
598 + const __m512 neg_quant_max_vec = _mm512_set1_ps(neg_quant_max);
599 +
600 + at::parallel_for(0, size, 4096, [&](int64_t start, int64_t end) {
601 + int64_t block_size = end - start;
602 + const float* block_src = src + start;
603 + float* block_dst = dst + start;
604 +
605 + int64_t i = 0;
606 + for (; i <= block_size - 16; i += 16) {
607 + __m512 src_vec = _mm512_loadu_ps(&block_src[i]);
608 + __m512 scaled_vec = _mm512_mul_ps(src_vec, scale_recip_vec);
609 + __m512 clamped_vec = _mm512_min_ps(_mm512_max_ps(scaled_vec, neg_quant_max_vec), quant_max_vec);
610 + _mm512_storeu_ps(&block_dst[i], clamped_vec);
611 + }
612 +
613 + // Handle remaining elements in block
614 + for (; i < block_size; ++i) {
615 + float scaled = block_src[i] * scale_reciprocal;
616 + block_dst[i] = std::clamp(scaled, neg_quant_max, quant_max);
617 + }
618 + });
619 +}
620 +
621 +std::tuple<at::Tensor, at::Tensor>
622 +_quantize_fp8e4m3(const at::Tensor& t, bool channelwise, c10::optional<at::Tensor> scale_opt = c10::nullopt) {
623 + constexpr float quant_max = 448.0f; // torch.finfo(torch.float8_e4m3fn).max
624 + constexpr float eps = std::numeric_limits<float>::epsilon();
625 +
626 + // Ensure input is contiguous and in float32
627 + auto t_float = t.to(at::ScalarType::Float).contiguous();
628 + at::Tensor qt;
629 + at::Tensor scale_tensor;
630 +
631 + if (channelwise) {
632 + // Channel-wise quantization with AVX512 optimization
633 + int64_t num_channels = t_float.size(0);
634 + int64_t elements_per_channel = t_float.numel() / num_channels;
635 +
636 + if (elements_per_channel * num_channels != t_float.numel()) {
637 + throw std::runtime_error("Tensor must be divisible by number of channels for channel-wise quantization");
638 + }
639 +
640 + // Allocate scale tensor
641 + scale_tensor = at::empty({num_channels}, t_float.options());
642 + float* scale_data = scale_tensor.data_ptr<float>();
643 +
644 + // Compute channel-wise max using AVX512
645 + compute_channel_max_avx512(t_float.data_ptr<float>(), scale_data, num_channels, elements_per_channel);
646 +
647 + // Apply quant_max and EPS
648 + at::parallel_for(0, num_channels, 1, [&](int64_t start, int64_t end) {
649 + for (int64_t c = start; c < end; ++c) {
650 + scale_data[c] = std::max(scale_data[c] / quant_max, eps);
651 + }
652 + });
653 +
654 + // Create output tensor for quantized values
655 + qt = at::empty_like(t_float);
656 +
657 + // Scale and clamp using AVX512
658 + scale_clamp_channelwise_avx512(
659 + t_float.data_ptr<float>(),
660 + qt.data_ptr<float>(),
661 + scale_data,
662 + num_channels,
663 + elements_per_channel,
664 + quant_max,
665 + -quant_max);
666 + } else {
667 + // Global quantization with AVX512 optimization
668 + if (!scale_opt.has_value()) {
669 + throw std::runtime_error("Scale must be provided for non-channelwise quantization in AVX512 version");
670 + }
671 +
672 + scale_tensor = scale_opt.value().to(at::ScalarType::Float).contiguous();
673 +
674 + // Handle scalar scale case
675 + float scale_val;
676 + if (scale_tensor.numel() == 1) {
677 + scale_val = std::max(scale_tensor.item<float>(), eps);
678 + } else {
679 + // Compute max of scale tensor if it's not scalar
680 + float* scale_data = scale_tensor.data_ptr<float>();
681 + __m512 max_vec = _mm512_set1_ps(eps);
682 +
683 + int64_t i = 0;
684 + for (; i <= scale_tensor.numel() - 16; i += 16) {
685 + __m512 scale_vec = _mm512_loadu_ps(&scale_data[i]);
686 + max_vec = _mm512_max_ps(max_vec, scale_vec);
687 + }
688 +
689 + float hmax = _mm512_reduce_max_ps(max_vec);
690 + for (; i < scale_tensor.numel(); ++i) {
691 + hmax = std::max(hmax, scale_data[i]);
692 + }
693 + scale_val = std::max(hmax, eps);
694 + }
695 +
696 + scale_tensor = at::tensor({scale_val}, scale_tensor.options());
697 + float scale_reciprocal = 1.0f / scale_val;
698 +
699 + // Create output tensor
700 + qt = at::empty_like(t_float);
701 +
702 + // Scale and clamp using AVX512
703 + scale_clamp_global_avx512(
704 + t_float.data_ptr<float>(), qt.data_ptr<float>(), t_float.numel(), scale_reciprocal, quant_max, -quant_max);
705 + }
706 +
707 + // Final conversion to FP8 E4M3FN using ATen's native conversion
708 + qt = at::_to_copy(qt, at::ScalarType::Float8_e4m3fn);
709 +
710 + return std::make_tuple(qt, scale_tensor);
711 +}
712 +
713 +inline __m128i cvtfp32_fp8e4m3(__m512& src) {
714 + // cvt 16x32 from fp32 to fp8 e4m3
715 + const __m512i sign_mask = _mm512_set1_epi32(0x80000000);
716 + const __m512i fp8_max = _mm512_set1_epi32(UINT32_C(1087) << 20);
717 + const __m512i denorm_thresh = _mm512_set1_epi32(UINT32_C(121) << 23);
718 + const __m512i denorm_mask = _mm512_set1_epi32(UINT32_C(141) << 23);
719 + const __m512i bias_part1 = _mm512_set1_epi32((uint32_t)(7 - 127) << 23);
720 + const __m512i rounding_bias = _mm512_set1_epi32(0x7FFFF);
721 + __m512i f_bits = _mm512_castps_si512(src);
722 + // Extract and save sign
723 + __m512i sign = _mm512_and_epi32(f_bits, sign_mask);
724 + f_bits = _mm512_xor_epi32(f_bits, sign);
725 +
726 + // Prepare result containers
727 + __m512i result = _mm512_setzero_si512();
728 +
729 + // Step 1: Handle case of overflow
730 + // (f_bits >= fp8_max): set result = 0x7f
731 + __mmask16 overflow_mask = _mm512_cmpge_epu32_mask(f_bits, fp8_max);
732 + if (overflow_mask) {
733 + result = _mm512_mask_set1_epi32(result, overflow_mask, 0x7f);
734 + }
735 +
736 + // Step 2: Handle small numbers (denormals)
737 + // Small numbers (f_bits < denorm_thresh)
738 + __mmask16 denorm_thresh_mask = _mm512_cmplt_epu32_mask(f_bits, denorm_thresh);
739 +
740 + if (denorm_thresh_mask) {
741 + __m512 small_input = _mm512_castsi512_ps(f_bits);
742 + __m512 small_denorm = _mm512_add_ps(small_input, _mm512_castsi512_ps(denorm_mask));
743 + __m512i small_denorm_bits = _mm512_castps_si512(small_denorm);
744 + __m512i small_result = _mm512_sub_epi32(small_denorm_bits, denorm_mask);
745 + result = _mm512_mask_mov_epi32(result, denorm_thresh_mask, small_result);
746 + }
747 +
748 + // Step 3: Handle normal numbers
749 + __mmask16 normal_mask = ~(overflow_mask | denorm_thresh_mask);
750 +
751 + if (normal_mask) {
752 + // mant_odd = (f_bits >> 20) & 1
753 + __m512i mant_odd = _mm512_and_epi32(_mm512_srli_epi32(f_bits, 20), _mm512_set1_epi32(1));
754 + // f_bits += bias_part1 + rounding_bias
755 + __m512i rounded = _mm512_add_epi32(f_bits, bias_part1);
756 + rounded = _mm512_add_epi32(rounded, rounding_bias);
757 + // Add mant_odd
758 + rounded = _mm512_add_epi32(rounded, mant_odd);
759 + // Shift right by 20 bits
760 + __m512i normal_result = _mm512_srli_epi32(rounded, 20);
761 + result = _mm512_mask_mov_epi32(result, normal_mask, normal_result);
762 + }
763 +
764 + // Merge back the sign
765 + __m512i sign_shifted = _mm512_srli_epi32(sign, 24);
766 + result = _mm512_or_epi32(result, sign_shifted);
767 +
768 + // Now result is 16 x 32-bit integers, but we only need 8-bit for each
769 + __m512i packed = _mm512_and_si512(result, _mm512_set1_epi32(0xFF));
770 +
771 + // Narrow 32-bit integers to 8-bit
772 + return _mm512_cvtepi32_epi8(packed);
773 +}
774 +
775 +__attribute__((target("avx10.2"))) inline __m128i cvtfp32_fp8e4m3_avx10_2(__m512& src) {
776 + __m256i f16_vec = _mm512_cvt_roundps_ph(src, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC);
777 + return _mm256_cvtph_hf8(_mm256_castsi256_ph(f16_vec));
778 +}
779 +
780 +// Carries the whole loop rather than just the conversion, as GCC cannot inline an avx10.2
781 +// function into a caller lacking that target, so branching per vector is costly.
782 +__attribute__((target("avx10.2"))) void quantize_bf16_block_avx10_2(
783 + const at::BFloat16* __restrict__ src, at::Float8_e4m3fn* __restrict__ dst, int64_t count, float scale_reciprocal) {
784 + constexpr float quant_max = 448.0f; // torch.finfo(torch.float8_e4m3fn).max
785 + const __m512 scale_recip_vec = _mm512_set1_ps(scale_reciprocal);
786 + const __m512 quant_max_vec = _mm512_set1_ps(quant_max);
787 + const __m512 neg_quant_max_vec = _mm512_set1_ps(-quant_max);
788 +
789 + for (int64_t i = 0; i <= count - 32; i += 32) {
790 + __m256i src_vec1 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&src[i]));
791 + __m256i src_vec2 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&src[i + 16]));
792 +
793 + __m512 fp32_vec1 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(src_vec1), 16));
794 + __m512 fp32_vec2 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(src_vec2), 16));
795 +
796 + __m512 scaled_vec1 = _mm512_mul_ps(fp32_vec1, scale_recip_vec);
797 + __m512 clamped_vec1 = _mm512_min_ps(_mm512_max_ps(scaled_vec1, neg_quant_max_vec), quant_max_vec);
798 +
799 + __m512 scaled_vec2 = _mm512_mul_ps(fp32_vec2, scale_recip_vec);
800 + __m512 clamped_vec2 = _mm512_min_ps(_mm512_max_ps(scaled_vec2, neg_quant_max_vec), quant_max_vec);
801 +
802 + __m128i fp8_vec1 = cvtfp32_fp8e4m3_avx10_2(clamped_vec1);
803 + __m128i fp8_vec2 = cvtfp32_fp8e4m3_avx10_2(clamped_vec2);
804 + _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + i), fp8_vec1);
805 + _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + i + 16), fp8_vec2);
806 + }
807 +}
808 +
809 +std::tuple<at::Tensor, at::Tensor> _quantize_fp8e4m3_bf16_per_tensor_no_scale(const at::Tensor& t) {
810 + constexpr float quant_max = 448.0f; // torch.finfo(torch.float8_e4m3fn).max
811 + constexpr float eps = std::numeric_limits<float>::epsilon();
812 +
813 + // Input validation and preparation
814 + assert(t.scalar_type() == at::ScalarType::BFloat16);
815 + auto t_bf16 = t.contiguous();
816 + int64_t num_channels = t_bf16.size(0);
817 + int64_t elements_per_channel = t_bf16.numel() / num_channels;
818 + assert(elements_per_channel % 32 == 0); // do not consider tile currently
819 + at::Tensor quant_t = at::empty(t_bf16.sizes(), t_bf16.options().dtype(at::kFloat8_e4m3fn));
820 +
821 + // Allocate output tensors
822 + at::Tensor scale_tensor = at::empty({num_channels}, t_bf16.options().dtype(at::ScalarType::Float));
823 + float* scale_data = scale_tensor.data_ptr<float>();
824 +
825 + // Unified processing: compute max, apply scale, and quantize in single pass
826 + at::parallel_for(0, num_channels, 1, [&](int64_t start, int64_t end) {
827 + for (int64_t c = start; c < end; ++c) {
828 + const at::BFloat16* channel_src = t_bf16.data_ptr<at::BFloat16>() + c * elements_per_channel;
829 + at::Float8_e4m3fn* quant_dst = quant_t.data_ptr<at::Float8_e4m3fn>() + c * elements_per_channel;
830 +
831 + // Step 1: Compute channel-wise max using AVX512
832 + float channel_max = 0.0f;
833 + int64_t i = 0;
834 +
835 + // Process 32 elements at a time for max computation
836 + for (; i <= elements_per_channel - 32; i += 32) {
837 + // Load 32 BF16 values (2x 256-bit vectors)
838 + __m256i src_vec1 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&channel_src[i]));
839 + __m256i src_vec2 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&channel_src[i + 16]));
840 +
841 + // Convert BF16 to FP32
842 + __m512i fp32_int1 = _mm512_cvtepu16_epi32(src_vec1);
843 + __m512 fp32_vec1 = _mm512_castsi512_ps(_mm512_slli_epi32(fp32_int1, 16));
844 +
845 + __m512i fp32_int2 = _mm512_cvtepu16_epi32(src_vec2);
846 + __m512 fp32_vec2 = _mm512_castsi512_ps(_mm512_slli_epi32(fp32_int2, 16));
847 +
848 + // Compute absolute values
849 + __m512 abs_vec1 = _mm512_abs_ps(fp32_vec1);
850 + __m512 abs_vec2 = _mm512_abs_ps(fp32_vec2);
851 +
852 + // Find max in each vector
853 + float max1 = _mm512_reduce_max_ps(abs_vec1);
854 + float max2 = _mm512_reduce_max_ps(abs_vec2);
855 + channel_max = std::max(channel_max, std::max(max1, max2));
856 + }
857 +
858 + // Step 2: Apply quant_max and EPS to compute scale
859 + float scale_val = std::max(channel_max / quant_max, eps);
860 + float scale_reciprocal = 1.0f / scale_val;
861 + scale_data[c] = scale_val;
862 +
863 + // Step 3: Scale and clamp using AVX512 (reuse the same loop structure)
864 + if (avx10_2_available()) {
865 + quantize_bf16_block_avx10_2(channel_src, quant_dst, elements_per_channel, scale_reciprocal);
866 + continue;
867 + }
868 + i = 0;
869 + const __m512 scale_recip_vec = _mm512_set1_ps(scale_reciprocal);
870 + const __m512 quant_max_vec = _mm512_set1_ps(quant_max);
871 + const __m512 neg_quant_max_vec = _mm512_set1_ps(-quant_max);
872 +
873 + for (; i <= elements_per_channel - 32; i += 32) {
874 + // Load 32 BF16 values
875 + __m256i src_vec1 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&channel_src[i]));
876 + __m256i src_vec2 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&channel_src[i + 16]));
877 +
878 + // Convert BF16 to FP32
879 + __m512 fp32_vec1 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(src_vec1), 16));
880 + __m512 fp32_vec2 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(src_vec2), 16));
881 +
882 + // Scale and clamp
883 + __m512 scaled_vec1 = _mm512_mul_ps(fp32_vec1, scale_recip_vec);
884 + __m512 clamped_vec1 = _mm512_min_ps(_mm512_max_ps(scaled_vec1, neg_quant_max_vec), quant_max_vec);
885 +
886 + __m512 scaled_vec2 = _mm512_mul_ps(fp32_vec2, scale_recip_vec);
887 + __m512 clamped_vec2 = _mm512_min_ps(_mm512_max_ps(scaled_vec2, neg_quant_max_vec), quant_max_vec);
888 +
889 + __m128i fp8_vec1 = cvtfp32_fp8e4m3(clamped_vec1);
890 + __m128i fp8_vec2 = cvtfp32_fp8e4m3(clamped_vec2);
891 + _mm_storeu_si128(reinterpret_cast<__m128i*>(quant_dst + i), fp8_vec1);
892 + _mm_storeu_si128(reinterpret_cast<__m128i*>(quant_dst + i + 16), fp8_vec2);
893 + }
894 + }
895 + });
896 +
897 + return std::make_tuple(quant_t, scale_tensor);
898 +}
899 +
900 +std::tuple<at::Tensor, at::Tensor>
901 +_quantize_fp8e4m3_bf16_per_tensor_with_scale(const at::Tensor& t, at::Tensor& scale_tensor) {
902 + constexpr float quant_max = 448.0f; // torch.finfo(torch.float8_e4m3fn).max
903 + constexpr float eps = std::numeric_limits<float>::epsilon();
904 +
905 + // Input validation
906 + assert(t.scalar_type() == at::ScalarType::BFloat16);
907 +
908 + auto t_bf16 = t.contiguous();
909 + auto scale_tensor_contig = scale_tensor.contiguous().to(at::ScalarType::Float);
910 +
911 + // Get scale value (handle scalar or tensor)
912 + float scale_val;
913 + if (scale_tensor_contig.numel() == 1) {
914 + scale_val = scale_tensor_contig.item<float>();
915 + } else {
916 + // Take max if scale is a tensor (though should be scalar for global quantization)
917 + scale_val = scale_tensor_contig.max().item<float>();
918 + }
919 +
920 + // Apply EPS to scale
921 + scale_val = std::max(scale_val, eps);
922 + float scale_reciprocal = 1.0f / scale_val;
923 +
924 + at::Tensor scale_output = at::tensor({scale_val}, scale_tensor_contig.options());
925 +
926 + // Apply scale and clamp using AVX512
927 + const at::BFloat16* src_data = t_bf16.data_ptr<at::BFloat16>();
928 + at::Tensor quant_t = at::empty(t_bf16.sizes(), t_bf16.options().dtype(at::kFloat8_e4m3fn));
929 + int64_t total_elements = t_bf16.numel();
930 + const __m512 scale_recip_vec = _mm512_set1_ps(scale_reciprocal);
931 + const __m512 quant_max_vec = _mm512_set1_ps(quant_max);
932 + const __m512 neg_quant_max_vec = _mm512_set1_ps(-quant_max);
933 + int64_t num_channels = t_bf16.size(0);
934 + int64_t elements_per_channel = total_elements / num_channels;
935 + assert(elements_per_channel % 32 == 0); // do not consider tile currently
936 + // Process in parallel blocks
937 + at::parallel_for(0, num_channels, 1, [&](int64_t start, int64_t end) {
938 + for (int64_t c = start; c < end; ++c) {
939 + const at::BFloat16* src_data = t_bf16.data_ptr<at::BFloat16>() + c * elements_per_channel;
940 + at::Float8_e4m3fn* quant_t_data = quant_t.data_ptr<at::Float8_e4m3fn>() + c * elements_per_channel;
941 +
942 + if (avx10_2_available()) {
943 + quantize_bf16_block_avx10_2(src_data, quant_t_data, elements_per_channel, scale_reciprocal);
944 + continue;
945 + }
946 +
947 + int64_t i = 0;
948 + // Process 32 elements at a time using AVX512
949 + for (; i <= elements_per_channel - 32; i += 32) {
950 + // Load 32 BF16 values (2x 256-bit vectors)
951 + __m256i src_vec1 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&src_data[i]));
952 + __m256i src_vec2 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&src_data[i + 16]));
953 +
954 + // Convert BF16 to FP32
955 + __m512 fp32_vec1 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(src_vec1), 16));
956 + __m512 fp32_vec2 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(src_vec2), 16));
957 +
958 + // Scale and clamp
959 + __m512 scaled_vec1 = _mm512_mul_ps(fp32_vec1, scale_recip_vec);
960 + __m512 clamped_vec1 = _mm512_min_ps(_mm512_max_ps(scaled_vec1, neg_quant_max_vec), quant_max_vec);
961 +
962 + __m512 scaled_vec2 = _mm512_mul_ps(fp32_vec2, scale_recip_vec);
963 + __m512 clamped_vec2 = _mm512_min_ps(_mm512_max_ps(scaled_vec2, neg_quant_max_vec), quant_max_vec);
964 +
965 + __m128i fp8_vec1 = cvtfp32_fp8e4m3(clamped_vec1);
966 + __m128i fp8_vec2 = cvtfp32_fp8e4m3(clamped_vec2);
967 + _mm_storeu_si128(reinterpret_cast<__m128i*>(quant_t_data + i), fp8_vec1);
968 + _mm_storeu_si128(reinterpret_cast<__m128i*>(quant_t_data + i + 16), fp8_vec2);
969 + }
970 + }
971 + });
972 +
973 + return std::make_tuple(quant_t, scale_output);
974 +}
975 +
976 +std::tuple<at::Tensor, at::Tensor>
977 +_quantize_fp8e4m3_vec(const at::Tensor& t, bool channelwise, c10::optional<at::Tensor> scale_opt) {
978 + if (channelwise) {
979 + return _quantize_fp8e4m3_bf16_per_tensor_no_scale(t);
980 + } else {
981 + assert(scale_opt.has_value());
982 + return _quantize_fp8e4m3_bf16_per_tensor_with_scale(t, scale_opt.value());
983 + }
984 +}
985 +
986 +// Public wrapper for torch op registration
987 +std::tuple<at::Tensor, at::Tensor>
988 +quantize_fp8e4m3_vec(const at::Tensor& t, bool channelwise, c10::optional<at::Tensor> scale_opt) {
989 + return _quantize_fp8e4m3_vec(t, channelwise, scale_opt);
990 +}
991 +
992 +at::Tensor fp8_scaled_mm_with_quant(
993 + const at::Tensor& act,
994 + const std::optional<at::Tensor>& act_scales,
995 + bool channelwise,
996 + const at::Tensor& weight,
997 + const at::Tensor& weight_scales,
998 + const std::optional<at::Tensor>& bias,
999 + at::ScalarType output_dtype) {
1000 + // TODO: refine here to use tensor ptr and in below dispatch api.
1001 + std::tuple<at::Tensor, at::Tensor> quant_act = act.scalar_type() == at::ScalarType::BFloat16
1002 + ? _quantize_fp8e4m3_vec(act, channelwise, act_scales)
1003 + : _quantize_fp8e4m3(act, channelwise, act_scales);
1004 + auto input = std::get<0>(quant_act);
1005 + auto input_scales = std::get<1>(quant_act);
1006 + int64_t N = weight.dim() == 4 ? weight.size(0) * weight.size(-1) : weight.size(0);
1007 + int act_quant_mode = input_scales.numel() == 1 ? PER_TENSOR
1008 + : input_scales.numel() == input.numel() / input.size(-1) ? PER_ROW
1009 + : PER_GROUP;
1010 + int wei_quant_mode = weight_scales.numel() == 1 ? PER_TENSOR : weight_scales.numel() == N ? PER_ROW : PER_GROUP;
1011 + // Case to fall back
1012 + if (weight.dim() == 2) {
1013 + TORCH_CHECK(
1014 + act_quant_mode != PER_GROUP && wei_quant_mode != PER_GROUP,
1015 + "FP8 linear: Per-group quantization is not supported in the fallback path");
1016 + auto y_fp32 = at::linear(input.to(at::kFloat).mul_(input_scales), weight.to(at::kFloat).mul_(weight_scales), bias);
1017 + return y_fp32.to(output_dtype);
1018 + }
1019 +
1020 + static bool cpublas_can_pack = cpublas_could_pack();
1021 + auto out_sizes = input.sizes().vec();
1022 + out_sizes.back() = N;
1023 + auto output = at::empty(out_sizes, input.options().dtype(output_dtype));
1024 +
1025 +#define AT_DISPATCH_FP8_LINEAR_KERNEL(OUT_DTYPE, CAN_PACK, A_QUANT_MODE, B_QUANT_MODE, ...) \
1026 + AT_DISPATCH_BOOL_NO_RETURN( \
1027 + CAN_PACK, \
1028 + "cpublas_can_pack", \
1029 + can_pack, \
1030 + AT_DISPATCH_QUANT_MODE_NO_RETURN( \
1031 + A_QUANT_MODE, \
1032 + "act_quant_mode", \
1033 + a_quant_mode, \
1034 + AT_DISPATCH_QUANT_MODE_NO_RETURN( \
1035 + B_QUANT_MODE, \
1036 + "wei_quant_mode", \
1037 + b_quant_mode, \
1038 + AT_DISPATCH_OUT_TYPES(OUT_DTYPE, "out_dtype", __VA_ARGS__))))
1039 +
1040 + AT_DISPATCH_FP8_LINEAR_KERNEL(output_dtype, cpublas_can_pack, act_quant_mode, wei_quant_mode, [&]() {
1041 + _float8_linear_impl<out_t, can_pack, a_quant_mode, b_quant_mode>(
1042 + input, input_scales, weight, weight_scales, bias, output);
1043 + });
1044 + return output;
1045 +}
1046 +
1047 +at::Tensor float8_linear_impl(
1048 + const at::Tensor& input,
1049 + const at::Tensor& input_scales,
1050 + const at::Tensor& weight,
1051 + const at::Tensor& weight_scales,
1052 + const std::optional<at::Tensor>& bias,
1053 + at::ScalarType output_dtype) {
1054 + int64_t N = weight.dim() == 4 ? weight.size(0) * weight.size(-1) : weight.size(0);
1055 + int act_quant_mode = input_scales.numel() == 1 ? PER_TENSOR
1056 + : input_scales.numel() == input.numel() / input.size(-1) ? PER_ROW
1057 + : PER_GROUP;
1058 + int wei_quant_mode = weight_scales.numel() == 1 ? PER_TENSOR : weight_scales.numel() == N ? PER_ROW : PER_GROUP;
1059 + // Case to fall back
1060 + if (weight.dim() == 2) {
1061 + TORCH_CHECK(
1062 + act_quant_mode != PER_GROUP && wei_quant_mode != PER_GROUP,
1063 + "FP8 linear: Per-group quantization is not supported in the fallback path");
1064 + auto y_fp32 = at::linear(input.to(at::kFloat).mul_(input_scales), weight.to(at::kFloat).mul_(weight_scales), bias);
1065 + return y_fp32.to(output_dtype);
1066 + }
1067 +
1068 + static bool cpublas_can_pack = cpublas_could_pack();
1069 + auto out_sizes = input.sizes().vec();
1070 + out_sizes.back() = N;
1071 + auto output = at::empty(out_sizes, input.options().dtype(output_dtype));
1072 +
1073 +#define AT_DISPATCH_FP8_LINEAR_KERNEL(OUT_DTYPE, CAN_PACK, A_QUANT_MODE, B_QUANT_MODE, ...) \
1074 + AT_DISPATCH_BOOL_NO_RETURN( \
1075 + CAN_PACK, \
1076 + "cpublas_can_pack", \
1077 + can_pack, \
1078 + AT_DISPATCH_QUANT_MODE_NO_RETURN( \
1079 + A_QUANT_MODE, \
1080 + "act_quant_mode", \
1081 + a_quant_mode, \
1082 + AT_DISPATCH_QUANT_MODE_NO_RETURN( \
1083 + B_QUANT_MODE, \
1084 + "wei_quant_mode", \
1085 + b_quant_mode, \
1086 + AT_DISPATCH_OUT_TYPES(OUT_DTYPE, "out_dtype", __VA_ARGS__))))
1087 +
1088 + AT_DISPATCH_FP8_LINEAR_KERNEL(output_dtype, cpublas_can_pack, act_quant_mode, wei_quant_mode, [&]() {
1089 + _float8_linear_impl<out_t, can_pack, a_quant_mode, b_quant_mode>(
1090 + input, input_scales, weight, weight_scales, bias, output);
1091 + });
1092 + return output;
1093 +}
\ No newline at end of file
csrc/cpu/sgl-kernels/moe.cpp
+93 -7
@@ -836,6 +836,11 @@ static inline void check_moe_scales(
836 TORCH_CHECK(w2_scale.has_value(), "missing w2_scale for mxfp4.");
837 TORCH_CHECK(w1_scale.value().scalar_type() == at::kByte, "expect w1_scale to be uint8.");
838 TORCH_CHECK(w2_scale.value().scalar_type() == at::kByte, "expect w2_scale to be uint8.");
839 + } else if constexpr (quant == CPUQuantMethod::FP8_W8A8) {
840 + TORCH_CHECK(w1_scale.has_value(), "missing w1_scale for fp8 w8a8.");
841 + TORCH_CHECK(w2_scale.has_value(), "missing w2_scale for fp8 w8a8.");
842 + TORCH_CHECK(block_size.has_value(), "missing block_size for fp8 w8a8.");
843 + TORCH_CHECK(block_size.value().size() == 2, "expect block_size.size() to be 2.");
844 }
845 }
846
@@ -850,9 +855,14 @@ static inline void check_moe_scales(
855 check_moe_scales<CPUQuantMethod::FP8_W8A16>(w1_scale, w2_scale, block_size);
856 } else if (moe_comp_method == CPUQuantMethod::MXFP4) {
857 check_moe_scales<CPUQuantMethod::MXFP4>(w1_scale, w2_scale, block_size);
858 + } else if (moe_comp_method == CPUQuantMethod::FP8_W8A8) {
859 + check_moe_scales<CPUQuantMethod::FP8_W8A8>(w1_scale, w2_scale, block_size);
860 }
861 }
862
863 +// CHECK_MOE_SCALES_FP8: validate scale shapes for FP8_W8A16.
864 +// w1s shape [E, 2N/block_N, K/block_K, block_N]: DIM0=1, DIM1=2
865 +// w2s shape [E, K/block_N, N/block_K]: DIM0=1, DIM1=2
866 #define CHECK_MOE_SCALES_FP8(DIM0, DIM1) \
867 auto w1s = w1_scale.value(); \
868 auto w2s = w2_scale.value(); \
@@ -864,6 +874,17 @@ static inline void check_moe_scales(
874 TORCH_CHECK(w2s.size(DIM0) == div_up(K, block_size_N)); \
875 TORCH_CHECK(w2s.size(DIM1) == div_up(N, block_size_K))
876
877 +// CHECK_MOE_SCALES_FP8_W8A8: extract scale info for FP8_W8A8 without shape
878 +// validation, because w1s is pre-packed by float8_linear_prepack_cpu with
879 +// shape [E, 2N/BLOCK_N, G, BLOCK_N] where BLOCK_N=32 != block_size_N(128).
880 +// This matches the sglang upstream behavior (checks are commented out there).
881 +#define CHECK_MOE_SCALES_FP8_W8A8(DIM0, DIM1) \
882 + auto w1s = w1_scale.value(); \
883 + auto w2s = w2_scale.value(); \
884 + auto block_size_val = block_size.value(); \
885 + int64_t block_size_N = block_size_val[0]; \
886 + int64_t block_size_K = block_size_val[1]
887 +
888 // hidden_states: [M, K]
889 // w1: [E, 2N, K] or [E, 2N, K / 2] for uint8
890 // w2: [E, K, N] or [E, K, N / 2] for uint8
@@ -883,6 +904,7 @@ void fused_experts_cpu(
904 const std::optional<at::Tensor>& w2_scale,
905 const std::optional<at::Tensor>& w1_zero,
906 const std::optional<at::Tensor>& w2_zero,
907 + const std::optional<at::Tensor>& a1_scale,
908 const std::optional<std::vector<int64_t>> block_size,
909 const std::optional<at::Tensor>& w1_bias,
910 const std::optional<at::Tensor>& w2_bias,
@@ -895,7 +917,10 @@ void fused_experts_cpu(
917 constexpr int64_t BLOCK_M = block_size_m();
918 constexpr int64_t BLOCK_N = block_size_n();
919
898 - const auto st = hidden_states.scalar_type();
920 + // FP8_W8A8 hidden_states input is FP8; scalar_t should be BF16 for output
921 + const auto st = (moe_comp_method == CPUQuantMethod::FP8_W8A8)
922 + ? at::kBFloat16
923 + : hidden_states.scalar_type();
924 // TODO: fused_topk_torch_native (CPU fallback for models like MiniMax)
925 // returns int64 topk_ids; fused_experts_cpu requires int32. Remove the typecast after topk kernel is provided
926 auto topk_ids_ = topk_ids.scalar_type() == at::kInt ? topk_ids : topk_ids.to(at::kInt);
@@ -908,6 +933,9 @@ void fused_experts_cpu(
933 if (moe_comp_method == CPUQuantMethod::INT4_W4A8 && is_vnni) {
934 CHECK_DIM(4, w1);
935 CHECK_DIM(4, w2);
936 + } else if (moe_comp_method == CPUQuantMethod::FP8_W8A8) {
937 + // w1 is prepackaged by float8_linear_prepack_cpu (5D), skip w1 dim check
938 + CHECK_DIM(3, w2);
939 } else {
940 CHECK_DIM(3, w1);
941 CHECK_DIM(3, w2);
@@ -924,8 +952,9 @@ void fused_experts_cpu(
952
953 int64_t M = hidden_states.size(0);
954 int64_t K = hidden_states.size(1);
927 - int64_t N = moe_comp_method == CPUQuantMethod::INT4_W4A8 ? w1_scale.value().size(1) * w1_scale.value().size(3) / 2
928 - : w1.size(1) / 2;
955 + int64_t N = moe_comp_method == CPUQuantMethod::INT4_W4A8
956 + ? w1_scale.value().size(1) * w1_scale.value().size(3) / 2
957 + : (moe_comp_method == CPUQuantMethod::FP8_W8A8 ? w2.size(2) : w1.size(1) / 2);
958 int64_t E = w1.size(0);
959 int64_t topk = topk_weights_.size(1);
960
@@ -935,7 +964,10 @@ void fused_experts_cpu(
964
965 // check weight shapes
966 CHECK_EQ(w2.size(0), E);
938 - if (!(moe_comp_method == CPUQuantMethod::INT4_W4A8)) {
967 + if (moe_comp_method == CPUQuantMethod::FP8_W8A8) {
968 + CHECK_EQ(w2.size(1), K);
969 + CHECK_EQ(packed_w2.size(2), packed_N);
970 + } else if (!(moe_comp_method == CPUQuantMethod::INT4_W4A8)) {
971 CHECK_EQ(w2.size(1), K);
972 CHECK_EQ(packed_w1.size(2), packed_K / (moe_comp_method == CPUQuantMethod::INT4_W4A8 ? 2 : 1));
973 CHECK_EQ(packed_w2.size(2), packed_N / (moe_comp_method == CPUQuantMethod::INT4_W4A8 ? 2 : 1));
@@ -993,22 +1025,35 @@ void fused_experts_cpu(
1025 // 5. Aq_tmp : [M, K] or [M * topk, N]
1026 // 6. As_tmp : [M * topk]
1027 //
996 - // for fp8 w8a16 and mxfp4:
1028 + // for fp8 w8a16, mxfp4, and fp8 w8a8:
1029 // 7. intermediate_cache0 : [M * topk, 2N]
1030 // 8. B_tmp : [T, MAX_CACHE_BLOCK_SIZE, BLOCK_N, std::max(K, N)]
1031 //
1032 + // for fp8 w8a8 additionally:
1033 + // 9. Ukernel_tmp : [T, 2 * BLOCK_M * BLOCK_N] (float32 for FP8×FP8 brgemm)
1034 + //
1035 int64_t buffer_size_nbytes =
1036 M * topk * N * 2 + M * topk * K * 2 +
1037 num_threads * BLOCK_M * K *
1003 - (moe_comp_method == CPUQuantMethod::INT8_W8A8 | moe_comp_method == CPUQuantMethod::INT4_W4A8 ? 1 : 2) +
1038 + (moe_comp_method == CPUQuantMethod::INT8_W8A8 | moe_comp_method == CPUQuantMethod::INT4_W4A8
1039 + | moe_comp_method == CPUQuantMethod::FP8_W8A8 ? 1 : 2) +
1040 num_threads * 2 * BLOCK_M * BLOCK_N * sizeof(float);
1041
1042 if (moe_comp_method == CPUQuantMethod::INT8_W8A8) {
1043 buffer_size_nbytes += std::max(M * K, M * topk * N) + M * topk * sizeof(float);
1044 }
1009 - if (moe_comp_method == CPUQuantMethod::FP8_W8A16 || moe_comp_method == CPUQuantMethod::MXFP4) {
1045 + if (moe_comp_method == CPUQuantMethod::FP8_W8A16 || moe_comp_method == CPUQuantMethod::MXFP4
1046 + || moe_comp_method == CPUQuantMethod::FP8_W8A8) {
1047 buffer_size_nbytes += M * topk * 2 * N * 2 + num_threads * MAX_CACHE_BLOCK_SIZE * BLOCK_N * std::max(K, N) * 2;
1048 }
1049 + if (moe_comp_method == CPUQuantMethod::FP8_W8A8) {
1050 + // Ukernel buffer: float32 workspace per thread for FP8×FP8 brgemm
1051 + buffer_size_nbytes += num_threads * 2 * BLOCK_M * BLOCK_N * sizeof(float);
1052 +#ifndef CPUBLAS_BRGEMM_F8F8F32
1053 + // BF16 fallback: also need dqA [T, BLOCK_M, BLOCK_K] + dqB [T, BLOCK_K, BLOCK_N]
1054 + buffer_size_nbytes += num_threads * (BLOCK_M * BLOCK_K + BLOCK_K * BLOCK_N) * 2;
1055 +#endif
1056 + }
1057 if (moe_comp_method == CPUQuantMethod::INT4_W4A8) {
1058 buffer_size_nbytes += M * topk * 2 * N * 2 + std::max(M * K, M * topk * N) + M * topk * sizeof(float) +
1059 num_threads * 2 * get_4bit_block_k_size(K / w1_scale.value().size(2)) * BLOCK_N;
@@ -1181,6 +1226,47 @@ void fused_experts_cpu(
1226 E,
1227 topk,
1228 num_tokens_post_pad);
1229 + } else if (moe_comp_method == CPUQuantMethod::FP8_W8A8) {
1230 + // Stage 1: FP8×FP8 GEMM (hidden_states is FP8 input)
1231 + // Stage 2: BF16×FP8 GEMM (W8A16 path for down-projection)
1232 + at::Float8_e4m3fn* __restrict__ A_tmp =
1233 + (at::Float8_e4m3fn*)((void*)(intermediate_cache2 + M * topk * K));
1234 + float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K));
1235 + scalar_t* __restrict__ intermediate_cache0 =
1236 + (scalar_t*)((void*)(C_tmp + num_threads * 2 * BLOCK_M * BLOCK_N));
1237 + scalar_t* __restrict__ B_tmp =
1238 + (scalar_t*)((void*)(intermediate_cache0 + M * topk * 2 * N));
1239 + float* __restrict__ Ukernel_tmp =
1240 + (float*)((void*)(B_tmp + num_threads * MAX_CACHE_BLOCK_SIZE * BLOCK_N * std::max(K, N)));
1241 +
1242 + CHECK_MOE_SCALES_FP8_W8A8(1, 2);
1243 + fused_experts_fp8_a8_kernel_impl(
1244 + out_hidden_states.data_ptr<scalar_t>(),
1245 + intermediate_cache0,
1246 + intermediate_cache1,
1247 + intermediate_cache2,
1248 + A_tmp,
1249 + B_tmp,
1250 + C_tmp,
1251 + Ukernel_tmp,
1252 + hidden_states.data_ptr<at::Float8_e4m3fn>(),
1253 + packed_w1.data_ptr<at::Float8_e4m3fn>(),
1254 + packed_w2.data_ptr<at::Float8_e4m3fn>(),
1255 + a1_scale.has_value() ? a1_scale.value().data_ptr<float>() : nullptr,
1256 + w1s.data_ptr<float>(),
1257 + w2s.data_ptr<float>(),
1258 + block_size_N,
1259 + block_size_K,
1260 + topk_weights_.data_ptr<float>(),
1261 + sorted_ids,
1262 + expert_ids,
1263 + offsets,
1264 + M,
1265 + N,
1266 + K,
1267 + E,
1268 + topk,
1269 + num_tokens_post_pad);
1270 } else {
1271 scalar_t* __restrict__ A_tmp = intermediate_cache2 + M * topk * K;
1272 float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K));
csrc/cpu/sgl-kernels/moe_fp8_w8a8.cpp new
+432
@@ -0,0 +1,432 @@
1 +// Adapted from
2 +// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
3 +
4 +// clang-format off
5 +
6 +#include <c10/util/Unroll.h>
7 +
8 +#include "common.h"
9 +#include "gemm.h"
10 +#include "moe.h"
11 +
12 +namespace {
13 +
14 +// Byte-wise copy for FP8 elements (no remainder assumption)
15 +inline void copy_stub2(
16 + at::Float8_e4m3fn* __restrict__ out,
17 + const at::Float8_e4m3fn* __restrict__ input,
18 + int64_t size) {
19 + for (int64_t d = 0; d < size; d++) {
20 + out[d] = input[d];
21 + }
22 +}
23 +
24 +#define PER_TENSOR 1
25 +#define PER_ROW 2
26 +#define PER_GROUP 3
27 +
28 +// Store float32 GEMM result to output buffer with dequantization and dtype
29 +// conversion. Activations are PER_ROW quantized; weights are PER_GROUP quantized.
30 +// Uses AVX512 on capable platforms.
31 +#if defined(CPU_CAPABILITY_AVX512)
32 +
33 +template <typename out_dtype, int64_t N, int act_quant_mode, int wei_quant_mode>
34 +inline void store_out(
35 + const float* y_buf,
36 + out_dtype* c_ptr,
37 + int64_t M,
38 + int64_t lda,
39 + const float* scales_a,
40 + const float* scales_b,
41 + const float* bias) {
42 + float a_scale = 1.0, b_scale = 1.0;
43 + __m512 va_scale, vb_scale;
44 + if constexpr (act_quant_mode == PER_TENSOR) {
45 + a_scale = *scales_a;
46 + }
47 + if constexpr (wei_quant_mode == PER_TENSOR) {
48 + b_scale = *scales_b;
49 + vb_scale = _mm512_set1_ps(b_scale);
50 + }
51 + for (int i = 0; i < M; ++i) {
52 + if constexpr (act_quant_mode == PER_ROW) {
53 + a_scale = *(scales_a + i);
54 + }
55 + if constexpr (act_quant_mode != PER_GROUP) {
56 + va_scale = _mm512_set1_ps(a_scale);
57 + }
58 + constexpr int N_UNROLL = N / 16;
59 + c10::ForcedUnroll<N_UNROLL>{}([&](auto idx) {
60 + constexpr int j = idx * 16;
61 + __m512 y_vec = _mm512_loadu_ps(y_buf + i * N + j);
62 + __m512 bias_vec = bias ? _mm512_loadu_ps(bias + j) : _mm512_setzero_ps();
63 + if constexpr (act_quant_mode != PER_GROUP) {
64 + y_vec = _mm512_mul_ps(y_vec, va_scale);
65 + }
66 + if constexpr (wei_quant_mode == PER_ROW) {
67 + vb_scale = _mm512_loadu_ps(scales_b + j);
68 + }
69 + if constexpr (wei_quant_mode != PER_GROUP) {
70 + y_vec = _mm512_mul_ps(y_vec, vb_scale);
71 + }
72 + y_vec = _mm512_add_ps(y_vec, bias_vec);
73 + if constexpr (std::is_same<out_dtype, float>::value) {
74 + _mm512_storeu_ps(c_ptr + i * lda + j, y_vec);
75 + } else if constexpr (std::is_same<out_dtype, at::BFloat16>::value) {
76 + __m256i y_bf16_vec = at::vec::cvtfp32_bf16(y_vec);
77 + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c_ptr + i * lda + j), y_bf16_vec);
78 + } else if constexpr (std::is_same<out_dtype, at::Half>::value) {
79 + __m256i y_fp16_vec = at::vec::cvtfp32_fp16(y_vec);
80 + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c_ptr + i * lda + j), y_fp16_vec);
81 + } else {
82 + TORCH_CHECK(false, "Unsupported output dtype");
83 + }
84 + });
85 + constexpr int tail_start = N / 16 * 16;
86 + for (int j = tail_start; j < N; ++j) {
87 + if constexpr (wei_quant_mode == PER_ROW) {
88 + b_scale = scales_b[j];
89 + }
90 + c_ptr[i * lda + j] = static_cast<out_dtype>(y_buf[i * N + j] * a_scale * b_scale);
91 + }
92 + } // for M
93 +}
94 +
95 +#else // no AVX512
96 +
97 +template <typename out_dtype, int64_t N, int act_quant_mode, int wei_quant_mode>
98 +inline void store_out(
99 + const float* y_buf,
100 + out_dtype* c_ptr,
101 + int64_t M,
102 + int64_t lda,
103 + const float* scales_a,
104 + const float* scales_b,
105 + const float* bias) {
106 + float a_scale = 1.0, b_scale = 1.0;
107 + if constexpr (act_quant_mode == PER_TENSOR) {
108 + a_scale = *scales_a;
109 + }
110 + if constexpr (wei_quant_mode == PER_TENSOR) {
111 + b_scale = *scales_b;
112 + }
113 + for (int i = 0; i < M; ++i) {
114 + if constexpr (act_quant_mode == PER_ROW) {
115 + a_scale = *(scales_a + i);
116 + }
117 + for (int j = 0; j < N; ++j) {
118 + if constexpr (wei_quant_mode == PER_ROW) {
119 + b_scale = scales_b[j];
120 + }
121 + float val = y_buf[i * N + j] * a_scale * b_scale;
122 + if (bias) val += bias[j];
123 + c_ptr[i * lda + j] = static_cast<out_dtype>(val);
124 + }
125 + }
126 +}
127 +
128 +#endif // CPU_CAPABILITY_AVX512
129 +
130 +} // anonymous namespace
131 +
132 +// FP8 W8A8 fused MoE kernel.
133 +//
134 +// hidden_states: [M, K] FP8 E4M3, pre-quantized activation
135 +// packed_w1: [E, 2*Nc, Kc, BLOCK_K, BLOCK_N] FP8 W8A8 prepack (float8_linear_prepack_cpu)
136 +// packed_w2: [E, K, N_packed] FP8 W8A16 VNNI pack (convert_weight_packed)
137 +// As: [M] float32 per-token activation scales
138 +// w1s: [E, num_groups, 2N] float32 per-group weight scales for w1
139 +// w2s: [E, ...] float32 per-group weight scales for w2
140 +//
141 +// Stage 1: hidden_states(FP8) @ w1(FP8) → ic0(BF16) [M*topk, 2N]
142 +// Uses FP8×FP8 GEMM (tinygemm_kernel with FP8 inputs)
143 +// Stage 1.5: ic1 = silu(ic0[:,:N]) * ic0[:,N:] [M*topk, N]
144 +// Stage 2: ic1(BF16) @ w2(FP8) → ic2(BF16) [M*topk, K]
145 +// Uses BF16×FP8 GEMM (W8A16 path, tinygemm_kernel<scalar_t>)
146 +// Stage 3: out = weighted_sum(ic2, topk_weights) [M, K]
147 +
148 +template <typename scalar_t>
149 +void fused_experts_fp8_a8_kernel_impl(
150 + scalar_t* __restrict__ output,
151 + scalar_t* __restrict__ ic0,
152 + scalar_t* __restrict__ ic1,
153 + scalar_t* __restrict__ ic2,
154 + at::Float8_e4m3fn* __restrict__ A_tmp,
155 + scalar_t* __restrict__ B_tmp,
156 + float* __restrict__ C_tmp,
157 + float* __restrict__ Ukernel_tmp,
158 + const at::Float8_e4m3fn* __restrict__ input,
159 + const at::Float8_e4m3fn* __restrict__ packed_w1,
160 + const at::Float8_e4m3fn* __restrict__ packed_w2,
161 + const float* __restrict__ As,
162 + const float* __restrict__ w1s,
163 + const float* __restrict__ w2s,
164 + int64_t block_size_N,
165 + int64_t block_size_K,
166 + const float* __restrict__ topk_weights,
167 + const int32_t* __restrict__ sorted_ids,
168 + const int32_t* __restrict__ expert_ids,
169 + const int32_t* __restrict__ offsets,
170 + int64_t M,
171 + int64_t N,
172 + int64_t K,
173 + int64_t E,
174 + int64_t topk,
175 + int64_t num_tokens_post_pad) {
176 + // 1. intermediate_cache1 : [M * topk, N]
177 + // 2. intermediate_cache2 : [M * topk, K]
178 + // 3. A_tmp : [T, BLOCK_M * K] FP8
179 + // 4. C_tmp : [T, 2 * BLOCK_M * BLOCK_N] float32
180 + // 5. intermediate_cache0 : [M * topk, 2N]
181 + // 6. B_tmp : [T, MAX_CACHE_BLOCK_SIZE, BLOCK_N, max(K, N)] scalar_t
182 + // 7. Ukernel_tmp: [T, 2 * BLOCK_M * BLOCK_N] float32
183 + // 8. (non-FP8-brgemm only) dqA [T, BLOCK_M * BLOCK_K] + dqB [T, BLOCK_K * BLOCK_N] in BF16
184 + // appended after Ukernel_tmp; see moe.cpp FP8_W8A8 buffer allocation
185 +
186 + constexpr int64_t BLOCK_M = block_size_m();
187 + constexpr int64_t BLOCK_N_VAL = block_size_n(); // = 32
188 + const int64_t KB = K / BLOCK_K; // K blocks
189 +
190 + int64_t B_tmp_size_per_thread = MAX_CACHE_BLOCK_SIZE * BLOCK_N_VAL * std::max(K, N);
191 +
192 + // stage 1: intermediate_cache0 = hidden_states @ w1
193 + const int64_t MB = div_up(num_tokens_post_pad, BLOCK_M);
194 + const int64_t NB = div_up(N, BLOCK_N_VAL);
195 +
196 + int64_t scale_size_N = div_up(2 * N, block_size_N);
197 + int64_t scale_size_K = div_up(K, block_size_K);
198 + int64_t blocks_n_per_group = block_size_N / BLOCK_N_VAL;
199 + int64_t num_groups = div_up(K, block_size_K); // G
200 + int64_t blocks_k_per_group = block_size_K / BLOCK_K;
201 + const int64_t stride_e = 2 * N * K; // per-expert stride for w1 (total elements)
202 + bool use_brgemm = true;
203 + int64_t num_thread = at::get_num_threads();
204 +
205 + // For the BF16 fallback path (no native FP8 brgemm), tinygemm_kernel needs temporary
206 + // BF16 buffers to dequantize FP8 inputs before running AMX-BF16 brgemm.
207 + // These are placed after the Ukernel_tmp region (see moe.cpp FP8_W8A8 buffer allocation).
208 +#ifndef CPUBLAS_BRGEMM_F8F8F32
209 + // dqA: [T, BLOCK_M, BLOCK_K] BF16
210 + // dqB: [T, BLOCK_K, BLOCK_N] BF16
211 + // Layout: immediately after Ukernel_tmp[num_thread * 2 * BLOCK_M * BLOCK_N_VAL]
212 + at::BFloat16* __restrict__ dq_base = reinterpret_cast<at::BFloat16*>(
213 + Ukernel_tmp + num_thread * 2 * BLOCK_M * BLOCK_N_VAL);
214 + const int64_t dq_per_thread = BLOCK_M * BLOCK_K + BLOCK_K * BLOCK_N_VAL;
215 +#endif
216 +
217 + at::parallel_for(0, MB * NB, 1, [&](int64_t begin, int64_t end) {
218 + int tid = get_thread_num();
219 + // A_tmp: FP8 activation buffer for this thread [BLOCK_M, K]
220 + at::Float8_e4m3fn* __restrict__ A = A_tmp + tid * BLOCK_M * K;
221 + float* __restrict__ C0 = C_tmp + tid * 2 * BLOCK_M * BLOCK_N_VAL;
222 + float* __restrict__ C1 = C0 + BLOCK_M * BLOCK_N_VAL;
223 + alignas(64) float As_[BLOCK_M];
224 + float* __restrict__ ukernel_buf_1 = Ukernel_tmp + tid * 2 * BLOCK_M * BLOCK_N_VAL;
225 +#ifndef CPUBLAS_BRGEMM_F8F8F32
226 + at::BFloat16* __restrict__ dqA_buf = dq_base + tid * dq_per_thread;
227 + at::BFloat16* __restrict__ dqB_buf = dqA_buf + BLOCK_M * BLOCK_K;
228 +#else
229 + at::BFloat16* __restrict__ dqA_buf = nullptr;
230 + at::BFloat16* __restrict__ dqB_buf = nullptr;
231 +#endif
232 + float* __restrict__ ukernel_buf_2 = ukernel_buf_1 + BLOCK_M * BLOCK_N_VAL;
233 + auto ldsa = 1; // activation per-row quantized
234 +
235 + for (int64_t i = begin; i < end; ++i) {
236 + int64_t mb = i / NB;
237 + int64_t nb = i % NB;
238 + int64_t nb1 = nb + NB; // gate is first N columns, up is second N columns
239 +
240 + int64_t n_size = std::min(N - nb * BLOCK_N_VAL, BLOCK_N_VAL);
241 +
242 + // B (weight) for expert expert_id: gate+up fused [2N, K] prepackaged as FP8
243 + int32_t expert_id = expert_ids[mb];
244 + const at::Float8_e4m3fn* __restrict__ B = packed_w1 + expert_id * stride_e;
245 + const float* __restrict__ Bs = w1s + expert_id * num_groups * (2 * N);
246 +
247 + // Load activation tile for this M-block
248 + const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
249 + int64_t m_size = offsets[mb + 1] - offsets[mb];
250 +
251 + for (int64_t m = 0; m < m_size; ++m) {
252 + int32_t index = A_ids[m] / topk;
253 + copy_stub2(A + m * K, input + index * K, K);
254 + As_[m] = As[index];
255 + }
256 +
257 + const int64_t offset = offsets[mb];
258 +
259 + // Gate path (nb): w1[:N, :]
260 + zero_buffer(C0, BLOCK_M * BLOCK_N_VAL);
261 + for (int kci = 0; kci < KB; ++kci) {
262 + // FP8×FP8 tinygemm: act_quant=PER_ROW(2), wei_quant=PER_GROUP(3)
263 + tinygemm_kernel(
264 + /* C */ C0,
265 + /* A */ A + kci * BLOCK_K,
266 + /* scales_a*/ As_ + kci / blocks_k_per_group,
267 + /* B */ B + (nb * KB + kci) * BLOCK_K * BLOCK_N_VAL,
268 + /* scales_b*/ Bs + nb * BLOCK_N_VAL * num_groups + kci / blocks_k_per_group * BLOCK_N_VAL,
269 + /* M */ m_size,
270 + /* K */ BLOCK_K,
271 + /* lda */ K,
272 + /* ldc */ BLOCK_N_VAL,
273 + /* ldsa */ ldsa,
274 + /* ukernel */ ukernel_buf_1,
275 + /* dqA_buf */ dqA_buf,
276 + /* dqB_buf */ dqB_buf);
277 + }
278 + store_out<scalar_t, BLOCK_N_VAL, PER_ROW, PER_GROUP>(
279 + C0,
280 + ic0 + offset * 2 * N + nb * BLOCK_N_VAL,
281 + m_size,
282 + 2 * N, // lda
283 + As_,
284 + nullptr,
285 + nullptr);
286 +
287 + // Up path (nb1): w1[N:, :]
288 + zero_buffer(C1, BLOCK_M * BLOCK_N_VAL);
289 + for (int kci = 0; kci < KB; ++kci) {
290 + tinygemm_kernel(
291 + /* C */ C1,
292 + /* A */ A + kci * BLOCK_K,
293 + /* scales_a*/ As_ + kci / blocks_k_per_group,
294 + /* B */ B + (nb1 * KB + kci) * BLOCK_K * BLOCK_N_VAL,
295 + /* scales_b*/ Bs + nb1 * BLOCK_N_VAL * num_groups + kci / blocks_k_per_group * BLOCK_N_VAL,
296 + /* M */ m_size,
297 + /* K */ BLOCK_K,
298 + /* lda */ K,
299 + /* ldc */ BLOCK_N_VAL,
300 + /* ldsa */ ldsa,
301 + /* ukernel */ ukernel_buf_2,
302 + /* dqA_buf */ dqA_buf,
303 + /* dqB_buf */ dqB_buf);
304 + }
305 + store_out<scalar_t, BLOCK_N_VAL, PER_ROW, PER_GROUP>(
306 + C1,
307 + ic0 + offset * 2 * N + nb1 * BLOCK_N_VAL,
308 + m_size,
309 + 2 * N, // lda
310 + As_,
311 + nullptr,
312 + nullptr);
313 + }
314 + if (use_brgemm) {
315 + at::native::cpublas::brgemm_release();
316 + }
317 + });
318 +
319 + // stage 1.5: intermediate_cache1 = silu(ic0[:, :N]) * ic0[:, N:]
320 + at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
321 + for (int64_t m = begin; m < end; ++m) {
322 + silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N);
323 + }
324 + });
325 +
326 + // stage 2: intermediate_cache2 = intermediate_cache1 @ w2
327 + // w2: [E, K, N] FP8, VNNI packed; treated as W8A16 (BF16 act × FP8 weight)
328 + const int64_t OC = K; // output channels (hidden dim)
329 + const int64_t IC = N; // input channels (expert intermediate dim)
330 + const int64_t MB2 = MB;
331 + const int64_t NB2 = div_up(OC, BLOCK_N_VAL);
332 + scale_size_N = div_up(K, block_size_N);
333 + scale_size_K = div_up(N, block_size_K);
334 + const int64_t stride_e2 = OC * IC; // per-expert stride for w2
335 + const int64_t stride_oc = IC;
336 +
337 + int64_t avg_M = std::max(int64_t(1), M * topk / E);
338 + bool use_brgemm2 = can_use_brgemm<at::Float8_e4m3fn>(avg_M);
339 +
340 + parallel_2d(MB2, NB2, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
341 + int tid = get_thread_num();
342 + alignas(64) scalar_t C[BLOCK_M * BLOCK_K];
343 +
344 + loop_2d<at::Float8_e4m3fn>(mb0, mb1, nb0, nb1, BLOCK_N_VAL * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
345 + int64_t m_size = offsets[mb + 1] - offsets[mb];
346 + int64_t n_size = std::min(OC - nb * BLOCK_N_VAL, BLOCK_N_VAL);
347 +
348 + // A from ic1: [M*topk, N] in sorted order (contiguous after stage 1.5)
349 + const scalar_t* __restrict__ A2 = ic1 + offsets[mb] * N;
350 + const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
351 +
352 + // B from w2: [E, K, N] VNNI packed, offset to expert + output block
353 + int32_t expert_id = expert_ids[mb];
354 + const at::Float8_e4m3fn* __restrict__ B2 =
355 + packed_w2 + expert_id * stride_e2 + nb * BLOCK_N_VAL * stride_oc;
356 + const float* __restrict__ Bs2 =
357 + w2s + expert_id * scale_size_N * scale_size_K + (nb / blocks_n_per_group) * scale_size_K;
358 +
359 + // Only unpack B on first visit or new expert
360 + int32_t pre_expert_id = mb == 0 ? -1 : expert_ids[mb - 1];
361 + bool do_unpack = (mb == mb0) || (expert_id != pre_expert_id);
362 +
363 + tinygemm_kernel<scalar_t>(
364 + /* A */ A2,
365 + /* B */ B2,
366 + /* C */ C,
367 + /* Btmp */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N_VAL * IC,
368 + /* Ctmp */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N_VAL,
369 + /* Bbias */ nullptr,
370 + /* scale */ Bs2,
371 + /* M */ m_size,
372 + /* N */ n_size,
373 + /* K */ IC,
374 + /* lda */ IC,
375 + /* ldb */ n_size,
376 + /* ldc */ BLOCK_N_VAL,
377 + /* brg */ use_brgemm2,
378 + /* block_sz_K */ block_size_K,
379 + /* do_unpack */ do_unpack);
380 +
381 + // Scatter ic2 in original token order, weighting by topk_weights
382 + for (int64_t m = 0; m < m_size; ++m) {
383 + int32_t index = A_ids[m];
384 + float weight = topk_weights[index];
385 + copy_mul_stub(ic2 + index * K + nb * BLOCK_N_VAL, C + m * BLOCK_N_VAL, weight, n_size);
386 + }
387 + });
388 +
389 + if (use_brgemm2) {
390 + at::native::cpublas::brgemm_release();
391 + }
392 + });
393 +
394 + // stage 3: output = sum(ic2, dim=topk) [M, K]
395 + at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
396 + for (int64_t m = begin; m < end; ++m) {
397 + sum_stub(output + m * K, ic2 + m * topk * K, topk, K);
398 + }
399 + });
400 +}
401 +
402 +#define INSTANTIATE_MOE_FP8_A8_TEMPLATE(TYPE) \
403 + template void fused_experts_fp8_a8_kernel_impl<TYPE>( \
404 + TYPE* __restrict__ output, \
405 + TYPE* __restrict__ ic0, \
406 + TYPE* __restrict__ ic1, \
407 + TYPE* __restrict__ ic2, \
408 + at::Float8_e4m3fn* __restrict__ A_tmp, \
409 + TYPE* __restrict__ B_tmp, \
410 + float* __restrict__ C_tmp, \
411 + float* __restrict__ Ukernel_tmp, \
412 + const at::Float8_e4m3fn* __restrict__ input, \
413 + const at::Float8_e4m3fn* __restrict__ packed_w1, \
414 + const at::Float8_e4m3fn* __restrict__ packed_w2, \
415 + const float* __restrict__ As, \
416 + const float* __restrict__ w1s, \
417 + const float* __restrict__ w2s, \
418 + int64_t block_size_N, \
419 + int64_t block_size_K, \
420 + const float* __restrict__ topk_weights, \
421 + const int32_t* __restrict__ sorted_ids, \
422 + const int32_t* __restrict__ expert_ids, \
423 + const int32_t* __restrict__ offsets, \
424 + int64_t M, \
425 + int64_t N, \
426 + int64_t K, \
427 + int64_t E, \
428 + int64_t topk, \
429 + int64_t num_tokens_post_pad)
430 +
431 +INSTANTIATE_MOE_FP8_A8_TEMPLATE(at::BFloat16);
432 +INSTANTIATE_MOE_FP8_A8_TEMPLATE(at::Half);
csrc/cpu/torch_bindings.cpp
+35 -2
@@ -73,6 +73,7 @@ void fused_experts_cpu(at::Tensor& out, at::Tensor& hidden_states,
73 const std::optional<at::Tensor>& w2_scale,
74 const std::optional<at::Tensor>& w1_zero,
75 const std::optional<at::Tensor>& w2_zero,
76 + const std::optional<at::Tensor>& a1_scale,
77 const std::optional<std::vector<int64_t>> block_size,
78 const std::optional<at::Tensor>& w1_bias,
79 const std::optional<at::Tensor>& w2_bias,
@@ -91,6 +92,18 @@ at::Tensor fp8_scaled_mm_cpu(at::Tensor& mat1, at::Tensor& mat2,
92 const std::optional<at::Tensor>& bias,
93 at::ScalarType out_dtype, bool is_vnni);
94
95 +// Adapted from sglang: FP8 W8A8 kernels
96 +std::tuple<at::Tensor, at::Tensor> float8_linear_prepack_impl(
97 + const at::Tensor& weight, const at::Tensor& scales);
98 +at::Tensor fp8_scaled_mm_with_quant(const at::Tensor& act,
99 + const std::optional<at::Tensor>& act_scales,
100 + bool channelwise, const at::Tensor& weight,
101 + const at::Tensor& weight_scales,
102 + const std::optional<at::Tensor>& bias,
103 + at::ScalarType output_dtype);
104 +std::tuple<at::Tensor, at::Tensor> quantize_fp8e4m3_vec(
105 + const at::Tensor& t, bool channelwise, c10::optional<at::Tensor> scale_opt);
106 +
107 // Adapted from sglang: MLA CPU kernels (AMX-only)
108 void decode_attention_cpu(at::Tensor& query, at::Tensor& k_buffer,
109 at::Tensor& v_buffer, at::Tensor& output,
@@ -281,6 +294,7 @@ void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,
294 const std::string& activation);
295
296 bool cpu_attn_has_isa(const std::string& isa);
297 +bool cpu_has_amx_fp8();
298
299 torch::Tensor get_scheduler_metadata(
300 const int64_t num_req, const int64_t num_heads_q,
@@ -611,7 +625,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
625 "fused_experts_cpu(Tensor(a0!) out, Tensor hidden_states, Tensor w1, "
626 "Tensor w2, Tensor topk_weights, Tensor topk_ids, "
627 "int moe_comp_method, Tensor? w1_scale, Tensor? w2_scale, "
614 - "Tensor? w1_zero, Tensor? w2_zero, int[]? block_size, "
628 + "Tensor? w1_zero, Tensor? w2_zero, Tensor? a1_scale, int[]? block_size, "
629 "Tensor? w1_bias, Tensor? w2_bias, float? alpha, float? limit, "
630 "bool is_vnni) -> "
631 "()");
@@ -629,7 +643,23 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
643 "bool is_vnni) -> Tensor");
644 ops.impl("fp8_scaled_mm_cpu", torch::kCPU, &fp8_scaled_mm_cpu);
645
632 - // Adapted from sglang: casual_conv1d kernels
646 + // Adapted from sglang: FP8 W8A8 kernels
647 + ops.def(
648 + "float8_linear_prepack_cpu(Tensor weight, Tensor scales) -> (Tensor, "
649 + "Tensor)");
650 + ops.impl("float8_linear_prepack_cpu", torch::kCPU,
651 + &float8_linear_prepack_impl);
652 + ops.def(
653 + "fp8_scaled_mm_with_quant(Tensor act, Tensor? act_scales, bool "
654 + "channelwise, "
655 + "Tensor weight, Tensor weight_scales, Tensor? bias, ScalarType "
656 + "output_dtype) -> Tensor");
657 + ops.impl("fp8_scaled_mm_with_quant", torch::kCPU, &fp8_scaled_mm_with_quant);
658 + ops.def(
659 + "quantize_fp8e4m3_vec(Tensor input, bool channelwise, Tensor? scale_opt) "
660 + "-> (Tensor, Tensor)");
661 + ops.impl("quantize_fp8e4m3_vec", torch::kCPU, &quantize_fp8e4m3_vec);
662 +
663 ops.def("causal_conv1d_weight_pack(Tensor weight) -> Tensor");
664 ops.impl("causal_conv1d_weight_pack", torch::kCPU,
665 &causal_conv1d_weight_pack);
@@ -832,6 +862,9 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
862 "-> (Tensor, Tensor)");
863 ops.impl("fused_gdn_gating_cpu", torch::kCPU, &fused_gdn_gating_cpu);
864
865 + // Pure runtime CPU ISA capability check.
866 + ops.def("cpu_has_amx_fp8() -> bool", &cpu_has_amx_fp8);
867 +
868 // CPU attention kernels
869 ops.def("cpu_attn_has_isa(str isa) -> bool", &cpu_attn_has_isa);
870 ops.def(
tests/kernels/moe/test_cpu_quant_fused_moe.py
+198
@@ -1114,5 +1114,203 @@ def test_int8_w8a8_cpu_fused_moe(M, N, K, E, topk, seed, is_vnni, inplace):
1114 )
1115
1116
1117 +# ===========================================================================
1118 +# FP8 W8A8 MoE
1119 +# ===========================================================================
1120 +
1121 +
1122 +requires_cpu_fp8_w8a8 = pytest.mark.skipif(
1123 + not ops.cpu_has_amx_fp8(),
1124 + reason="requires native AMX-FP8 hardware",
1125 +)
1126 +
1127 +FP8_W8A8_MAX = torch.finfo(torch.float8_e4m3fn).max
1128 +FP8_W8A8_QUANT_GROUP = 128 # quantization group size for both K and N dimensions
1129 +
1130 +
1131 +def _make_fp8_w8a8_weight_w13(
1132 + E: int, two_n: int, K: int, group_K: int = FP8_W8A8_QUANT_GROUP
1133 +) -> tuple[torch.Tensor, torch.Tensor]:
1134 + """Create w13 [E, 2N, K] FP8 with group scales [E, 2N, G] where G = K//group_K."""
1135 + G = K // group_K
1136 + w_list, s_list = [], []
1137 + for _ in range(E):
1138 + w_f32 = torch.randn(two_n, K)
1139 + # Per-row, per-K-group quantization: scale [2N, G]
1140 + w_re = w_f32.view(two_n, G, group_K)
1141 + abs_max = w_re.abs().amax(dim=2, keepdim=True).clamp(min=1e-7) # [2N, G, 1]
1142 + scale = (abs_max / FP8_W8A8_MAX).squeeze(2) # [2N, G]
1143 + w_q = (
1144 + (w_re / abs_max).clamp(-FP8_W8A8_MAX, FP8_W8A8_MAX).to(torch.float8_e4m3fn)
1145 + )
1146 + w_list.append(w_q.view(two_n, K).contiguous())
1147 + s_list.append(scale.float())
1148 + return torch.stack(w_list), torch.stack(s_list) # [E, 2N, K], [E, 2N, G]
1149 +
1150 +
1151 +def _make_fp8_w8a8_weight_w2(
1152 + E: int,
1153 + K: int,
1154 + N: int,
1155 + group_K: int = FP8_W8A8_QUANT_GROUP,
1156 + group_N: int = FP8_W8A8_QUANT_GROUP,
1157 +) -> tuple[torch.Tensor, torch.Tensor]:
1158 + """Create w2 [E, K, N] FP8 with block scales [E, K//gK, N//gN]."""
1159 + nK = K // group_K
1160 + nN = N // group_N
1161 + w_list, s_list = [], []
1162 + for _ in range(E):
1163 + w_f32 = torch.randn(K, N)
1164 + w_re = w_f32.view(nK, group_K, nN, group_N)
1165 + abs_max = w_re.abs().amax(dim=(1, 3), keepdim=True).clamp(min=1e-7)
1166 + scale = (abs_max / FP8_W8A8_MAX).squeeze(1).squeeze(2) # [nK, nN]
1167 + w_q = (
1168 + (w_re / abs_max).clamp(-FP8_W8A8_MAX, FP8_W8A8_MAX).to(torch.float8_e4m3fn)
1169 + )
1170 + w_list.append(w_q.view(K, N).contiguous())
1171 + s_list.append(scale.float())
1172 + return torch.stack(w_list), torch.stack(s_list) # [E, K, N], [E, K//gK, N//gN]
1173 +
1174 +
1175 +def _pack_fp8_w8a8_w13(
1176 + w13: torch.Tensor, # [E, 2N, K] FP8
1177 + w13_scale: torch.Tensor, # [E, 2N, G] float32
1178 +) -> tuple[torch.Tensor, torch.Tensor]:
1179 + """Pack w13 for CPU FP8 W8A8 MoE kernel using float8_linear_prepack_cpu."""
1180 + E = w13.size(0)
1181 + packed_list, scale_list = [], []
1182 + for i in range(E):
1183 + pw, ps = torch.ops._C.float8_linear_prepack_cpu(
1184 + w13[i].contiguous(), w13_scale[i].contiguous()
1185 + )
1186 + packed_list.append(pw)
1187 + scale_list.append(ps)
1188 + return torch.stack(packed_list), torch.stack(scale_list)
1189 +
1190 +
1191 +def _ref_fp8_w8a8_moe(
1192 + hidden_states: torch.Tensor, # BF16 [M, K_hidden] (after FP8 round-trip dequant)
1193 + w13: torch.Tensor, # FP8 [E, 2N, K_hidden]
1194 + w13_scale: torch.Tensor, # float32 [E, 2N, G] (G = K_hidden//group_K)
1195 + w2: torch.Tensor, # FP8 [E, K_hidden, N]
1196 + w2_scale: torch.Tensor, # float32 [E, K_hidden//gK, N//gN]
1197 + topk_weights: torch.Tensor, # float32 [M, top_k]
1198 + topk_ids: torch.Tensor, # int64 [M, top_k]
1199 + N: int,
1200 + group_K: int = FP8_W8A8_QUANT_GROUP,
1201 + group_N: int = FP8_W8A8_QUANT_GROUP,
1202 +) -> torch.Tensor:
1203 + """Reference implementation: dequantize weights then compute MoE."""
1204 + M, K_hidden = hidden_states.shape
1205 + top_k = topk_ids.shape[1]
1206 + output = torch.zeros(M, K_hidden, dtype=torch.bfloat16)
1207 +
1208 + for tok in range(M):
1209 + x = hidden_states[tok].float() # [K_hidden]
1210 + for k_idx in range(top_k):
1211 + expert_id = topk_ids[tok, k_idx].item()
1212 + w = topk_weights[tok, k_idx].item()
1213 +
1214 + # Stage 1: x @ w13.T — gate and up projections
1215 + w13_e = w13[expert_id].float() # [2N, K_hidden]
1216 + ws_e = w13_scale[expert_id].float() # [2N, G]
1217 + G = ws_e.shape[1]
1218 + w13_dq = torch.zeros_like(w13_e)
1219 + for g in range(G):
1220 + c0, c1 = g * group_K, (g + 1) * group_K
1221 + w13_dq[:, c0:c1] = w13_e[:, c0:c1] * ws_e[:, g : g + 1]
1222 + gate_up = (x.unsqueeze(0) @ w13_dq.T).squeeze(0) # [2N]
1223 +
1224 + gate = gate_up[:N]
1225 + up = gate_up[N:]
1226 + act = F.silu(gate) * up # [N]
1227 +
1228 + # Stage 2: act @ w2.T — down projection
1229 + w2_e = w2[expert_id].float() # [K_hidden, N]
1230 + ws2_e = w2_scale[expert_id].float() # [K_hidden//gK, N//gN]
1231 + nK2 = K_hidden // group_K
1232 + nN2 = N // group_N
1233 + w2_dq = torch.zeros_like(w2_e)
1234 + for gi in range(nK2):
1235 + for gj in range(nN2):
1236 + r0, r1 = gi * group_K, (gi + 1) * group_K
1237 + c0, c1 = gj * group_N, (gj + 1) * group_N
1238 + w2_dq[r0:r1, c0:c1] = w2_e[r0:r1, c0:c1] * ws2_e[gi, gj]
1239 + out_e = (act.unsqueeze(0) @ w2_dq.T).squeeze(0) # [K_hidden]
1240 +
1241 + output[tok] += (w * out_e).to(torch.bfloat16)
1242 +
1243 + return output
1244 +
1245 +
1246 +@requires_cpu_fp8_w8a8
1247 +@pytest.mark.parametrize(
1248 + "E,N,K,M,top_k",
1249 + [
1250 + # N and K must be multiples of FP8_W8A8_QUANT_GROUP=128
1251 + (4, 128, 256, 1, 2), # single token (decode phase)
1252 + (4, 128, 256, 8, 2),
1253 + (8, 256, 512, 16, 2),
1254 + (4, 256, 512, 4, 1),
1255 + ],
1256 +)
1257 +def test_fp8_w8a8_cpu_fused_moe(E: int, N: int, K: int, M: int, top_k: int):
1258 + """Test fused_experts_cpu FP8_W8A8 shape/dtype and accuracy vs a
1259 + dequantized BF16 reference."""
1260 + set_random_seed(42)
1261 +
1262 + w13, w13_scale = _make_fp8_w8a8_weight_w13(E, 2 * N, K)
1263 + w2, w2_scale = _make_fp8_w8a8_weight_w2(E, K, N)
1264 +
1265 + packed_w13, packed_w13_scale = _pack_fp8_w8a8_w13(w13, w13_scale)
1266 + packed_w2 = _prepack_experts(w2)
1267 +
1268 + hidden_states = torch.randn(M, K, dtype=torch.bfloat16) * 0.1
1269 + x_fp8, x_scales = torch.ops._C.quantize_fp8e4m3_vec(hidden_states, True, None)
1270 + assert x_fp8.dtype == torch.float8_e4m3fn
1271 + assert x_scales.shape == (M,)
1272 +
1273 + gen = torch.Generator().manual_seed(0)
1274 + topk_weights = torch.softmax(torch.randn(M, top_k, generator=gen), dim=-1).float()
1275 + topk_ids = torch.zeros(M, top_k, dtype=torch.int32)
1276 + for i in range(M):
1277 + topk_ids[i] = torch.randperm(E, generator=torch.Generator().manual_seed(i))[
1278 + :top_k
1279 + ].int()
1280 +
1281 + output = torch.empty(M, K, dtype=torch.bfloat16)
1282 + ops.fused_experts_cpu(
1283 + output,
1284 + x_fp8,
1285 + packed_w13,
1286 + packed_w2,
1287 + topk_weights,
1288 + topk_ids,
1289 + ops.CPUQuantMethod.FP8_W8A8,
1290 + packed_w13_scale, # w1_scale: [E, Nc, G, BLOCK_N]
1291 + w2_scale, # w2_scale: [E, K//128, N//128]
1292 + None, # w1_zero
1293 + None, # w2_zero
1294 + [FP8_W8A8_QUANT_GROUP, FP8_W8A8_QUANT_GROUP], # block_size
1295 + None, # w1_bias
1296 + None, # w2_bias
1297 + None, # alpha
1298 + None, # limit
1299 + True, # is_vnni (weights already packed)
1300 + x_scales, # a1_scale: per-token FP8 activation scales
1301 + )
1302 +
1303 + assert output.shape == (M, K)
1304 + assert output.dtype == torch.bfloat16
1305 + assert not torch.isnan(output).any()
1306 + assert not torch.isinf(output).any()
1307 +
1308 + x_dequant = (x_fp8.float() * x_scales.unsqueeze(1)).bfloat16()
1309 + ref = _ref_fp8_w8a8_moe(
1310 + x_dequant, w13, w13_scale, w2, w2_scale, topk_weights, topk_ids.long(), N
1311 + )
1312 + torch.testing.assert_close(output.float(), ref.float(), atol=1e-2, rtol=1e-2)
1313 +
1314 +
1315 if __name__ == "__main__":
1316 sys.exit(pytest.main([__file__, "-v"]))
tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
+225 -4
@@ -1,6 +1,10 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Tests for CPU FP8 W8A16 block-scaled GEMM kernel (fp8_scaled_mm_cpu).
3 +"""Tests for CPU FP8 scaled-mm GEMM kernels.
4 +
5 +Covers:
6 + * fp8_scaled_mm_cpu (W8A16)
7 + * fp8_scaled_mm_with_quant (W8A8)
8
9 Run `pytest tests/kernels/quantization/test_cpu_fp8_scaled_mm.py -v`.
10 """
@@ -23,19 +27,31 @@ from vllm.platforms import current_platform
27 if not current_platform.is_cpu():
28 pytest.skip("skipping CPU-only tests", allow_module_level=True)
29
26 -if not ops._supports_cpu_fp8_w8a16:
27 - pytest.skip("fp8_scaled_mm_cpu op not available", allow_module_level=True)
30 +requires_cpu_fp8_w8a16 = pytest.mark.skipif(
31 + not ops._supports_cpu_fp8_w8a16, reason="fp8_scaled_mm_cpu op not available"
32 +)
33 +requires_cpu_fp8_w8a8 = pytest.mark.skipif(
34 + not ops.cpu_has_amx_fp8(),
35 + reason="requires native AMX-FP8 hardware",
36 +)
37
38 if not torch.cpu._is_amx_tile_supported():
39 pytest.skip("requires AMX tile support", allow_module_level=True)
40
41 +FP8_MAX = torch.finfo(torch.float8_e4m3fn).max
42 BLOCK_SIZE = [128, 128]
43 +GROUP_SIZE = 128
44
45
46 def cdiv(a: int, b: int) -> int:
47 return -(a // -b)
48
49
50 +# ---------------------------------------------------------------------------
51 +# W8A16 helpers
52 +# ---------------------------------------------------------------------------
53 +
54 +
55 def quantize_weight_block_fp8(
56 weight: torch.Tensor,
57 block_size: list[int],
@@ -126,7 +142,7 @@ def ref_fp8_block_scaled_mm(
142
143
144 # ---------------------------------------------------------------------------
129 -# Test parameters
145 +# W8A16 test: fp8_scaled_mm_cpu
146 # ---------------------------------------------------------------------------
147 M_SIZES = [1, 4, 16, 64, 128]
148 # (N, K) — weight shape is [N, K], output has N columns.
@@ -141,6 +157,7 @@ NK_SIZES = [
157 ]
158
159
160 +@requires_cpu_fp8_w8a16
161 @pytest.mark.parametrize("M", M_SIZES)
162 @pytest.mark.parametrize("N,K", NK_SIZES)
163 @pytest.mark.parametrize("use_bias", [False, True])
@@ -200,6 +217,77 @@ def ref_fp8_per_tensor_scaled_mm(
217 return out.to(out_dtype)
218
219
220 +# ---------------------------------------------------------------------------
221 +# W8A8 helpers
222 +# ---------------------------------------------------------------------------
223 +
224 +
225 +def quantize_weight_per_tensor(
226 + weight: torch.Tensor,
227 +) -> tuple[torch.Tensor, torch.Tensor]:
228 + """Quantize weight [N, K] to FP8 with a single per-tensor scale, broadcast
229 + to [N, 1] (the shape float8_linear_prepack_cpu expects)."""
230 + N = weight.shape[0]
231 + abs_max = weight.abs().max()
232 + scale = (abs_max / FP8_MAX).clamp(min=1e-7)
233 + q = (weight.float() / scale).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn)
234 + return q, scale.expand(N).contiguous().view(N, 1)
235 +
236 +
237 +def quantize_weight_per_channel(
238 + weight: torch.Tensor,
239 +) -> tuple[torch.Tensor, torch.Tensor]:
240 + """Quantize weight [N, K] to FP8 with per-channel (per-row) scales [N, 1]."""
241 + abs_max = weight.abs().amax(dim=1, keepdim=True)
242 + scale = (abs_max / FP8_MAX).clamp(min=1e-7)
243 + q = (weight.float() / scale).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn)
244 + return q, scale.float()
245 +
246 +
247 +def quantize_weight_per_group(
248 + weight: torch.Tensor,
249 + group_size: int = GROUP_SIZE,
250 +) -> tuple[torch.Tensor, torch.Tensor]:
251 + """Quantize weight [N, K] to FP8 with per-group (128-block) scales [N, G]."""
252 + N, K = weight.shape
253 + G = K // group_size
254 + w_groups = weight.view(N, G, group_size)
255 + abs_max = w_groups.abs().amax(dim=2, keepdim=True)
256 + scale = (abs_max / FP8_MAX).clamp(min=1e-7)
257 + q = (w_groups.float() / scale).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn)
258 + return q.view(N, K), scale.squeeze(2).float()
259 +
260 +
261 +def quantize_act_per_token(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
262 + """Reference per-token (row) dynamic FP8 quantization of activations."""
263 + x_f = x.float()
264 + scale = (x_f.abs().amax(dim=1, keepdim=True) / FP8_MAX).clamp(min=1e-7)
265 + scaled = (x_f / scale).clamp(-FP8_MAX, FP8_MAX)
266 + # Round via fp16 to match the AVX10.2 kernel path.
267 + q = scaled.half().to(torch.float8_e4m3fn)
268 + return q, scale
269 +
270 +
271 +def reference_static_fp8_linear(
272 + x: torch.Tensor, # BF16 [M, K]
273 + weight_fp8: torch.Tensor, # FP8 [N, K]
274 + weight_scale: torch.Tensor, # float32 [N, 1]
275 + act_scale: torch.Tensor, # float32 scalar
276 + bias: torch.Tensor | None,
277 + out_dtype: torch.dtype,
278 +) -> torch.Tensor:
279 + """Static per-tensor: quantize/dequant x with act_scale, dequant weight, matmul."""
280 + x_scaled = (x.float() / act_scale.item()).clamp(-FP8_MAX, FP8_MAX)
281 + # Round via fp16 to match the AVX10.2 kernel path.
282 + x_q = x_scaled.half().to(torch.float8_e4m3fn)
283 + x_dq = x_q.float() * act_scale.item()
284 + w_dq = weight_fp8.float() * weight_scale.view(-1, 1)
285 + out = torch.mm(x_dq, w_dq.t())
286 + if bias is not None:
287 + out = out + bias.float()
288 + return out.to(out_dtype)
289 +
290 +
291 NK_SIZES_PER_TENSOR = [
292 (32, 64),
293 (5120, 5120),
@@ -274,3 +362,136 @@ def test_cpu_fp8_per_tensor_kernel_rejects_non_multiple_of_32_n(n: int):
362 )
363 supported, _ = CPUFp8PerTensorScaledMMLinearKernel.can_implement(config)
364 assert not supported
365 +
366 +
367 +def reference_dynamic_fp8_linear(
368 + x: torch.Tensor, # BF16 [M, K]
369 + weight_fp8: torch.Tensor, # FP8 [N, K]
370 + weight_scale: torch.Tensor, # float32 [N, G]
371 + bias: torch.Tensor | None,
372 + out_dtype: torch.dtype,
373 +) -> torch.Tensor:
374 + """Dynamic per-token act x per-group weight: quantize x per-token,
375 + dequant weight per-group, matmul."""
376 + N, K = weight_fp8.shape
377 + G = weight_scale.shape[1]
378 + group_size = K // G
379 + x_q, x_scale = quantize_act_per_token(x)
380 + x_dq = x_q.float() * x_scale
381 + w_dq = (
382 + weight_fp8.float().view(N, G, group_size) * weight_scale.view(N, G, 1)
383 + ).view(N, K)
384 + out = torch.mm(x_dq, w_dq.t())
385 + if bias is not None:
386 + out = out + bias.float()
387 + return out.to(out_dtype)
388 +
389 +
390 +# ---------------------------------------------------------------------------
391 +# W8A8 tests: fp8_scaled_mm_with_quant
392 +# ---------------------------------------------------------------------------
393 +W8A8_M_SIZES = [1, 4, 16, 64]
394 +W8A8_NK_SIZES = [(128, 256), (256, 512), (64, 128)]
395 +
396 +
397 +@requires_cpu_fp8_w8a8
398 +@pytest.mark.parametrize("static_activation", [True, False])
399 +@pytest.mark.parametrize("M", W8A8_M_SIZES)
400 +@pytest.mark.parametrize("N,K", W8A8_NK_SIZES)
401 +@pytest.mark.parametrize("use_bias", [False, True])
402 +def test_cpu_fp8_w8a8_per_tensor(
403 + M: int, N: int, K: int, use_bias: bool, static_activation: bool
404 +):
405 + """CPUFP8W8A8ScaledMMLinearKernel with per-tensor weight scale."""
406 + torch.manual_seed(42)
407 + out_dtype = torch.bfloat16
408 +
409 + x = torch.randn(M, K, dtype=torch.bfloat16)
410 + weight = torch.randn(N, K, dtype=torch.bfloat16)
411 + bias = torch.randn(N, dtype=torch.float32) if use_bias else None
412 +
413 + w_fp8, w_scale = quantize_weight_per_tensor(weight)
414 + packed_weight, packed_scale = torch.ops._C.float8_linear_prepack_cpu(w_fp8, w_scale)
415 +
416 + if static_activation:
417 + act_scale = (x.float().abs().max() / FP8_MAX).clamp(min=1e-7).view(1)
418 + kernel_out = ops.fp8_scaled_mm_with_quant(
419 + x, act_scale, False, packed_weight, packed_scale, bias, out_dtype
420 + )
421 + ref_out = reference_static_fp8_linear(
422 + x, w_fp8, w_scale, act_scale, bias, out_dtype
423 + )
424 + atol = 0.05
425 + else:
426 + kernel_out = ops.fp8_scaled_mm_with_quant(
427 + x, None, True, packed_weight, packed_scale, bias, out_dtype
428 + )
429 + ref_out = reference_dynamic_fp8_linear(x, w_fp8, w_scale, bias, out_dtype)
430 + atol = 0.05
431 +
432 + assert kernel_out.dtype == out_dtype
433 + torch.testing.assert_close(kernel_out, ref_out, rtol=0.05, atol=atol)
434 +
435 +
436 +@requires_cpu_fp8_w8a8
437 +@pytest.mark.parametrize("static_activation", [True, False])
438 +@pytest.mark.parametrize("M", W8A8_M_SIZES)
439 +@pytest.mark.parametrize("N,K", W8A8_NK_SIZES)
440 +@pytest.mark.parametrize("use_bias", [False, True])
441 +def test_cpu_fp8_w8a8_per_channel(
442 + M: int, N: int, K: int, use_bias: bool, static_activation: bool
443 +):
444 + """CPUFP8W8A8ScaledMMLinearKernel with per-channel weight scale."""
445 + torch.manual_seed(42)
446 + out_dtype = torch.bfloat16
447 +
448 + x = torch.randn(M, K, dtype=torch.bfloat16)
449 + weight = torch.randn(N, K, dtype=torch.bfloat16)
450 + bias = torch.randn(N, dtype=torch.float32) if use_bias else None
451 +
452 + w_fp8, w_scale = quantize_weight_per_channel(weight)
453 + packed_weight, packed_scale = torch.ops._C.float8_linear_prepack_cpu(w_fp8, w_scale)
454 +
455 + if static_activation:
456 + act_scale = (x.float().abs().max() / FP8_MAX).clamp(min=1e-7).view(1)
457 + kernel_out = ops.fp8_scaled_mm_with_quant(
458 + x, act_scale, False, packed_weight, packed_scale, bias, out_dtype
459 + )
460 + ref_out = reference_static_fp8_linear(
461 + x, w_fp8, w_scale, act_scale, bias, out_dtype
462 + )
463 + atol = 0.05
464 + else:
465 + kernel_out = ops.fp8_scaled_mm_with_quant(
466 + x, None, True, packed_weight, packed_scale, bias, out_dtype
467 + )
468 + ref_out = reference_dynamic_fp8_linear(x, w_fp8, w_scale, bias, out_dtype)
469 + atol = 0.05
470 +
471 + assert kernel_out.dtype == out_dtype
472 + torch.testing.assert_close(kernel_out, ref_out, rtol=0.05, atol=atol)
473 +
474 +
475 +@requires_cpu_fp8_w8a8
476 +@pytest.mark.parametrize("M", W8A8_M_SIZES)
477 +@pytest.mark.parametrize("N,K", W8A8_NK_SIZES)
478 +@pytest.mark.parametrize("use_bias", [False, True])
479 +def test_cpu_fp8_w8a8_per_group(M: int, N: int, K: int, use_bias: bool):
480 + """CPUFp8W8A8BlockScaledMMKernel: per-group (128-block) weight scale +
481 + dynamic per-token activation scale."""
482 + torch.manual_seed(42)
483 + out_dtype = torch.bfloat16
484 +
485 + x = torch.randn(M, K, dtype=torch.bfloat16)
486 + weight = torch.randn(N, K, dtype=torch.bfloat16)
487 + bias = torch.randn(N, dtype=torch.float32) if use_bias else None
488 +
489 + w_fp8, w_scale = quantize_weight_per_group(weight)
490 + packed_weight, packed_scale = torch.ops._C.float8_linear_prepack_cpu(w_fp8, w_scale)
491 + kernel_out = ops.fp8_scaled_mm_with_quant(
492 + x, None, True, packed_weight, packed_scale, bias, out_dtype
493 + )
494 + ref_out = reference_dynamic_fp8_linear(x, w_fp8, w_scale, bias, out_dtype)
495 +
496 + assert kernel_out.dtype == out_dtype
497 + torch.testing.assert_close(kernel_out, ref_out, rtol=0.05, atol=0.05)
vllm/_custom_ops.py
+44
@@ -2,6 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from enum import IntEnum
5 +from functools import cache
6 from typing import TYPE_CHECKING, Literal
7
8 import torch
@@ -3384,6 +3385,7 @@ class CPUQuantMethod(IntEnum):
3385 FP8_W8A16 = 2
3386 INT4_W4A8 = 3
3387 MXFP4 = 4
3388 + FP8_W8A8 = 5
3389
3390
3391 if hasattr(torch.ops._C, "dynamic_4bit_int_moe"):
@@ -3422,6 +3424,7 @@ def fused_experts_cpu(
3424 alpha: float | None = None,
3425 limit: float | None = None,
3426 is_vnni: bool = True,
3427 + a1_scale: torch.Tensor | None = None,
3428 ) -> None:
3429 torch.ops._C.fused_experts_cpu(
3430 out,
@@ -3435,6 +3438,7 @@ def fused_experts_cpu(
3438 w2_scale,
3439 w1_zero,
3440 w2_zero,
3441 + a1_scale,
3442 block_size,
3443 w1_bias,
3444 w2_bias,
@@ -3564,6 +3568,46 @@ def fp8_scaled_mm_cpu(
3568 )
3569
3570
3571 +# FP8 W8A8 CPU kernels
3572 +@cache
3573 +def cpu_has_amx_fp8() -> bool:
3574 + """Whether this CPU has native AMX-FP8 MMA support."""
3575 + if not hasattr(torch.ops._C, "cpu_has_amx_fp8"):
3576 + return False
3577 + return bool(torch.ops._C.cpu_has_amx_fp8())
3578 +
3579 +
3580 +if hasattr(torch.ops._C, "fp8_scaled_mm_with_quant"):
3581 +
3582 + @register_fake("_C::fp8_scaled_mm_with_quant")
3583 + def fp8_scaled_mm_with_quant_fake(
3584 + act: torch.Tensor,
3585 + act_scales: torch.Tensor | None,
3586 + channelwise: bool,
3587 + weight: torch.Tensor,
3588 + weight_scales: torch.Tensor,
3589 + bias: torch.Tensor | None,
3590 + output_dtype: torch.dtype,
3591 + ) -> torch.Tensor:
3592 + M = act.reshape(-1, act.size(-1)).size(0)
3593 + N = weight.size(0) * weight.size(-1)
3594 + return torch.empty((M, N), dtype=output_dtype, device=act.device)
3595 +
3596 +
3597 +def fp8_scaled_mm_with_quant(
3598 + act: torch.Tensor,
3599 + act_scales: torch.Tensor | None,
3600 + channelwise: bool,
3601 + weight: torch.Tensor,
3602 + weight_scales: torch.Tensor,
3603 + bias: torch.Tensor | None,
3604 + output_dtype: torch.dtype,
3605 +) -> torch.Tensor:
3606 + return torch.ops._C.fp8_scaled_mm_with_quant(
3607 + act, act_scales, channelwise, weight, weight_scales, bias, output_dtype
3608 + )
3609 +
3610 +
3611 def chunk_gated_delta_rule_cpu(
3612 query: torch.Tensor,
3613 key: torch.Tensor,
vllm/model_executor/kernels/linear/__init__.py
+4
@@ -186,6 +186,8 @@ from vllm.model_executor.kernels.linear.scaled_mm.b12x import (
186 from vllm.model_executor.kernels.linear.scaled_mm.cpu import (
187 CPUFp8BlockScaledMMKernel,
188 CPUFp8PerTensorScaledMMLinearKernel,
189 + CPUFp8W8A8BlockScaledMMKernel,
190 + CPUFP8W8A8ScaledMMLinearKernel,
191 CPUInt8ScaledMMLinearKernel,
192 )
193 from vllm.model_executor.kernels.linear.scaled_mm.cutlass import (
@@ -436,6 +438,7 @@ _POSSIBLE_FP8_KERNELS: dict[PlatformEnum, list[type[FP8ScaledMMLinearKernel]]] =
438 ChannelWiseTorchFP8ScaledMMLinearKernel,
439 ],
440 PlatformEnum.CPU: [
441 + CPUFP8W8A8ScaledMMLinearKernel,
442 CPUFp8PerTensorScaledMMLinearKernel,
443 PerTensorTorchFP8ScaledMMLinearKernel,
444 ChannelWiseTorchFP8ScaledMMLinearKernel,
@@ -468,6 +471,7 @@ _POSSIBLE_FP8_BLOCK_KERNELS: dict[
471 TritonFp8BlockScaledMMKernel,
472 ],
473 PlatformEnum.CPU: [
474 + CPUFp8W8A8BlockScaledMMKernel, # W8A8 preferred; falls back to W8A16 below
475 CPUFp8BlockScaledMMKernel,
476 ],
477 PlatformEnum.XPU: [
vllm/model_executor/kernels/linear/scaled_mm/__init__.py
+2
@@ -10,6 +10,7 @@ from vllm.model_executor.kernels.linear.scaled_mm.BlockScaledMMLinearKernel impo
10 from vllm.model_executor.kernels.linear.scaled_mm.cpu import (
11 CPUFp8BlockScaledMMKernel,
12 CPUFp8PerTensorScaledMMLinearKernel,
13 + CPUFP8W8A8ScaledMMLinearKernel,
14 CPUInt8ScaledMMLinearKernel,
15 )
16 from vllm.model_executor.kernels.linear.scaled_mm.cutlass import (
@@ -69,5 +70,6 @@ __all__ = [
70 "Fp8BlockScaledMMLinearKernel",
71 "CPUFp8BlockScaledMMKernel",
72 "CPUFp8PerTensorScaledMMLinearKernel",
73 + "CPUFP8W8A8ScaledMMLinearKernel",
74 "XPUFp8BlockScaledMMKernel",
75 ]
vllm/model_executor/kernels/linear/scaled_mm/cpu.py
+303 -7
@@ -270,12 +270,17 @@ class CPUFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
270 # Skip the base class process (FP8 padding / fnuz normalization)
271 # which is GPU-oriented. Instead, VNNI-prepack weights for AMX.
272 params = self._get_layer_params(layer)
273 - packed_weight = torch.ops._C.convert_weight_packed(params.weight)
274 - replace_parameter(
275 - layer,
276 - params.WEIGHT,
277 - torch.nn.Parameter(packed_weight, requires_grad=False),
278 - )
273 +
274 + # MLA's kv_b_proj (see `_cpu_skip_gemm_dispatch`): MLA reads this
275 + # weight to derive W_UK/W_UV, so keep the checkpoint [N, K] layout and
276 + # let the kernel pack on the fly instead.
277 + if not getattr(layer, "_cpu_skip_gemm_dispatch", False):
278 + packed_weight = torch.ops._C.convert_weight_packed(params.weight)
279 + replace_parameter(
280 + layer,
281 + params.WEIGHT,
282 + torch.nn.Parameter(packed_weight, requires_grad=False),
283 + )
284
285 # Re-wrap scale as a plain Parameter so the kernel can read it
286 # without weight-loader metadata interfering.
@@ -320,7 +325,7 @@ class CPUFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
325 list(self.weight_group_shape),
326 bias,
327 x.dtype,
323 - True, # is_vnni (weight already prepacked)
328 + not getattr(layer, "_cpu_skip_gemm_dispatch", False), # is_vnni
329 )
330 return out.reshape(x.shape[:-1] + (out.size(-1),)) if x.dim() > 2 else out
331
@@ -430,3 +435,294 @@ class CPUFp8PerTensorScaledMMLinearKernel(FP8ScaledMMLinearKernel):
435 raise NotImplementedError(
436 "CPUFp8PerTensorScaledMMLinearKernel overrides apply_weights directly."
437 )
438 +
439 +
440 +class CPUFp8W8A8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
441 + """FP8 W8A8 block-quantized GEMM with dynamic per-token activation quant on CPU."""
442 +
443 + # Activation quant is done inside fp8_scaled_mm_with_quant;
444 + # the base class must not apply a second quant step.
445 + apply_input_quant = False
446 +
447 + @classmethod
448 + def is_supported(
449 + cls, compute_capability: int | None = None
450 + ) -> tuple[bool, str | None]:
451 + if not current_platform.is_cpu():
452 + return False, "requires CPU platform."
453 + if not torch.cpu._is_amx_tile_supported():
454 + return False, "requires AMX tile support (Sapphire Rapids or newer)."
455 + if not ops.cpu_has_amx_fp8():
456 + return False, (
457 + "requires native AMX-FP8 hardware; falls back "
458 + "to CPUFp8BlockScaledMMKernel (W8A16) otherwise."
459 + )
460 + return True, None
461 +
462 + @classmethod
463 + def can_implement(
464 + cls, config: FP8ScaledMMLinearLayerConfig
465 + ) -> tuple[bool, str | None]:
466 + # Same block-shape requirements as the W8A16 sibling kernel.
467 + weight_gs = config.weight_quant_key.scale.group_shape
468 + if weight_gs.col <= 0 or weight_gs.col != 128:
469 + return False, (
470 + "CPU FP8 W8A8 block kernel requires K-dimension block size of 128, "
471 + f"got {weight_gs.col}."
472 + )
473 + if weight_gs.row <= 0 or weight_gs.row % 32 != 0:
474 + return False, (
475 + "CPU FP8 W8A8 block kernel requires N-dimension block size to be "
476 + f"a positive multiple of 32, got {weight_gs.row}."
477 + )
478 + if config.out_dtype not in (torch.bfloat16, torch.float32):
479 + return False, "Only bfloat16/float32 output dtype supported."
480 + n = config.weight_shape[0]
481 + if n % 32 != 0:
482 + # float8_linear_prepack_cpu tiles N in chunks of 32 and cannot
483 + # handle a remainder tile smaller than that.
484 + return False, f"requires weight output dim (N={n}) to be a multiple of 32."
485 + return True, None
486 +
487 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
488 + """Prepack block-quantized FP8 weight for W8A8 AMX BRGEMM.
489 +
490 + For block-quant FP8 models weight is stored as [N, K] (NOT transposed,
491 + unlike non-block FP8 where vLLM stores [K, N]).
492 +
493 + float8_linear_prepack_cpu expects:
494 + weight: [N, K]
495 + scales: [N, G] where G = K / block_k
496 +
497 + weight_scale_inv from the checkpoint has shape [ceil(N/block_n), G],
498 + so we expand rows from ceil(N/block_n) to N via repeat_interleave.
499 + """
500 + params = self._get_layer_params(layer)
501 + weight = params.weight # [N, K]
502 +
503 + # MLA's kv_b_proj (see `_cpu_skip_gemm_dispatch`): MLA reads this
504 + # weight to derive W_UK/W_UV, so keep the checkpoint [N, K] layout and
505 + # its block scales; apply_weights falls back to the W8A16 op.
506 + if getattr(layer, "_cpu_skip_gemm_dispatch", False):
507 + return
508 +
509 + weight_scale = (
510 + params.weight_scale_inv
511 + if params.weight_scale_inv is not None
512 + else params.weight_scale
513 + ) # [ceil(N/block_n), ceil(K/block_k)]
514 + assert weight_scale is not None, "FP8 W8A8 block kernel requires weight scale."
515 +
516 + N = weight.shape[0]
517 + block_n = self.weight_group_shape.row # e.g. 128
518 +
519 + # Expand scale rows: [ceil(N/block_n), G] → [N, G]
520 + if weight_scale.size(0) != N:
521 + weight_scale_expanded = weight_scale.repeat_interleave(block_n, dim=0)[
522 + :N, :
523 + ].contiguous()
524 + else:
525 + weight_scale_expanded = weight_scale.contiguous()
526 +
527 + packed_weight, packed_scales = torch.ops._C.float8_linear_prepack_cpu(
528 + weight, weight_scale_expanded.to(torch.float32)
529 + )
530 +
531 + scale_attr = (
532 + params.WEIGHT_SCALE_INV
533 + if params.weight_scale_inv is not None
534 + else params.WEIGHT_SCALE
535 + )
536 + replace_parameter(
537 + layer,
538 + params.WEIGHT,
539 + torch.nn.Parameter(packed_weight, requires_grad=False),
540 + )
541 + replace_parameter(
542 + layer,
543 + scale_attr,
544 + torch.nn.Parameter(packed_scales, requires_grad=False),
545 + )
546 +
547 + def apply_weights(
548 + self,
549 + layer: torch.nn.Module,
550 + x: torch.Tensor,
551 + bias: torch.Tensor | None = None,
552 + **kwargs,
553 + ) -> torch.Tensor:
554 + params = self._get_layer_params(layer)
555 + weight_scale = (
556 + params.weight_scale_inv
557 + if params.weight_scale_inv is not None
558 + else params.weight_scale
559 + )
560 +
561 + x_2d = x.reshape(-1, x.shape[-1]) if x.dim() > 2 else x
562 + if getattr(layer, "_cpu_skip_gemm_dispatch", False):
563 + # Unpacked weight: fp8_scaled_mm_with_quant has no is_vnni option,
564 + # so fall back to the W8A16 op (BF16 activation, block-scaled FP8
565 + # weight) for this layer.
566 + out = torch.ops._C.fp8_scaled_mm_cpu(
567 + x_2d,
568 + params.weight,
569 + weight_scale,
570 + list(self.weight_group_shape),
571 + bias,
572 + x.dtype,
573 + False, # is_vnni
574 + )
575 + return out.reshape(x.shape[:-1] + (out.size(-1),)) if x.dim() > 2 else out
576 +
577 + # act_scales=None + channelwise=True → dynamic per-token activation quant
578 + out = ops.fp8_scaled_mm_with_quant(
579 + x_2d,
580 + None,
581 + True,
582 + params.weight,
583 + weight_scale,
584 + bias,
585 + x.dtype,
586 + )
587 + return out.reshape(x.shape[:-1] + (out.size(-1),)) if x.dim() > 2 else out
588 +
589 + def apply_block_scaled_mm(
590 + self,
591 + A: torch.Tensor,
592 + B: torch.Tensor,
593 + As: torch.Tensor,
594 + Bs: torch.Tensor,
595 + ) -> torch.Tensor:
596 + raise NotImplementedError(
597 + "CPUFp8W8A8BlockScaledMMKernel overrides apply_weights directly."
598 + )
599 +
600 +
601 +class CPUFP8W8A8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
602 + """FP8 W8A8 GEMM with dynamic per-token activation quantization on CPU."""
603 +
604 + @classmethod
605 + def is_supported(
606 + cls, compute_capability: int | None = None
607 + ) -> tuple[bool, str | None]:
608 + if not current_platform.is_cpu():
609 + return False, "requires CPU platform."
610 + if not torch.cpu._is_amx_tile_supported():
611 + return False, "requires AMX tile support (Sapphire Rapids or newer)."
612 + if not ops.cpu_has_amx_fp8():
613 + return False, (
614 + "requires native AMX-FP8 hardware (Diamond Rapids); falls back "
615 + "to the Torch FP8 kernels otherwise."
616 + )
617 + return True, None
618 +
619 + @classmethod
620 + def can_implement(
621 + cls, config: FP8ScaledMMLinearLayerConfig
622 + ) -> tuple[bool, str | None]:
623 + if config.out_dtype not in (torch.bfloat16, torch.float16, torch.float32):
624 + return False, "Only bfloat16/float16/float32 output dtype supported."
625 + # Reject block-quantized weights — those go to CPUFp8BlockScaledMMKernel
626 + weight_gs = config.weight_quant_key.scale.group_shape
627 + if weight_gs.col > 0 and weight_gs.row > 0:
628 + return False, "Block-quantized weights not handled by FP8 W8A8 kernel."
629 + n = config.weight_shape[0]
630 + if n % 32 != 0:
631 + # float8_linear_prepack_cpu tiles N in chunks of 32 and cannot
632 + # handle a remainder tile smaller than that.
633 + return False, f"requires weight output dim (N={n}) to be a multiple of 32."
634 + return True, None
635 +
636 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
637 + """Prepack weights with float8_linear_prepack_cpu (VNNI + block layout)."""
638 + w_name, ws_name, xs_name, _ = self.layer_param_names
639 + weight = getattr(layer, w_name)
640 + weight_scale = getattr(layer, ws_name)
641 +
642 + # vLLM stores linear weight as [K, N], while the CPU prepack op expects [N, K].
643 + N = weight.shape[-1]
644 +
645 + # Ensure scale is 2D [N, G]
646 + ws = weight_scale
647 + if ws is None:
648 + raise RuntimeError("FP8 W8A8 kernel requires weight scale.")
649 + ws = ws.to(torch.float32)
650 +
651 + # Fused modules may carry per-shard scales (e.g. QKV/MLP shards).
652 + # Expand those to per-output-channel scales before prepack.
653 + if (
654 + len(getattr(layer, "logical_widths", [])) > 1
655 + and ws.dim() <= 1
656 + and ws.numel() == len(layer.logical_widths)
657 + ):
658 + ws = convert_to_channelwise(ws.reshape(-1), layer.logical_widths)
659 +
660 + if ws.dim() == 0:
661 + ws = ws.view(1, 1).expand(N, 1)
662 + elif ws.dim() == 1:
663 + if ws.numel() == 1:
664 + ws = ws.view(1, 1).expand(N, 1)
665 + elif ws.numel() == N:
666 + ws = ws.unsqueeze(1) # [N] -> [N, 1]
667 + else:
668 + raise RuntimeError(
669 + "FP8 W8A8 kernel got unsupported 1D weight scale shape: "
670 + f"{tuple(ws.shape)} for N={N}."
671 + )
672 +
673 + # Align scale rows to output channels N.
674 + if ws.size(0) != N:
675 + if ws.size(0) == 1:
676 + ws = ws.expand(N, ws.size(1)).contiguous()
677 + else:
678 + raise RuntimeError(
679 + "FP8 W8A8 kernel got mismatched weight scale rows: "
680 + f"ws_rows={ws.size(0)}, expected N={N}."
681 + )
682 +
683 + w_2d = weight.t().contiguous() # [N, K]
684 + packed_weight, packed_scales = torch.ops._C.float8_linear_prepack_cpu(w_2d, ws)
685 +
686 + replace_parameter(
687 + layer,
688 + w_name,
689 + torch.nn.Parameter(packed_weight, requires_grad=False),
690 + )
691 + replace_parameter(
692 + layer,
693 + ws_name,
694 + torch.nn.Parameter(packed_scales, requires_grad=False),
695 + )
696 +
697 + def apply_weights(
698 + self,
699 + layer: torch.nn.Module,
700 + x: torch.Tensor,
701 + bias: torch.Tensor | None = None,
702 + **kwargs,
703 + ) -> torch.Tensor:
704 + w_name, ws_name, xs_name, _ = self.layer_param_names
705 + weight = getattr(layer, w_name)
706 + weight_scale = getattr(layer, ws_name)
707 + # For static quantization, use pre-computed activation scale
708 + act_scales = getattr(layer, xs_name, None) if xs_name else None
709 +
710 + x_2d = x.reshape(-1, x.shape[-1]) if x.dim() > 2 else x
711 + # channelwise=True: per-token (per-row) quantization when act_scales is None
712 + # channelwise=False: per-tensor quantization when act_scales is provided
713 + channelwise = act_scales is None
714 + out = ops.fp8_scaled_mm_with_quant(
715 + x_2d,
716 + act_scales,
717 + channelwise,
718 + weight,
719 + weight_scale,
720 + bias,
721 + x.dtype,
722 + )
723 + return out.reshape(x.shape[:-1] + (out.size(-1),)) if x.dim() > 2 else out
724 +
725 + def apply_scaled_mm(self, A, B, out_dtype, As, Bs, bias=None, **kwargs):
726 + raise NotImplementedError(
727 + "CPUFP8W8A8ScaledMMLinearKernel uses apply_weights directly."
728 + )
vllm/model_executor/layers/attention/mla_attention.py
+2 -2
@@ -1055,7 +1055,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
1055 mqa_q_nope,
1056 self.impl._w_uk_packed, # type: ignore[attr-defined]
1057 True,
1058 - None,
1058 + self.impl._w_scale, # type: ignore[attr-defined]
1059 )
1060 mqa_ql_nope = mqa_ql_nope.transpose(0, 1)
1061 else:
@@ -1401,7 +1401,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
1401 x,
1402 self.impl._w_uv_packed, # type: ignore[attr-defined]
1403 True,
1404 - None,
1404 + self.impl._w_scale, # type: ignore[attr-defined]
1405 )
1406 else:
1407 # Multiply + Transpose (N, B, L) x (N, L, V)->(N, B, V)->(B, N, V)
vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
+217
@@ -15,6 +15,7 @@ from vllm._custom_ops import (
15 convert_weight_packed_scale_zp,
16 cpu_fused_moe,
17 cpu_fused_moe_int8,
18 + cpu_has_amx_fp8,
19 cpu_prepack_moe_weight,
20 cpu_prepack_moe_weight_int8,
21 fused_experts_cpu,
@@ -467,6 +468,7 @@ class CPUExpertsFp8(mk.FusedMoEExpertsModular):
468 RoutingMethodType.Default,
469 RoutingMethodType.Renormalize,
470 RoutingMethodType.RenormalizeNaive,
471 + RoutingMethodType.DeepSeekV3,
472 RoutingMethodType.DeepseekV4,
473 RoutingMethodType.DeepSeekV3,
474 ]
@@ -1612,3 +1614,218 @@ class ZenCPUExpertsInt8(mk.FusedMoEExpertsModular):
1614 self._w1_scale_bf16,
1615 self._w2_scale_bf16,
1616 )
1617 +
1618 +
1619 +# ===========================================================================
1620 +# FP8 W8A8 MoE
1621 +# ===========================================================================
1622 +
1623 +
1624 +def prepare_fp8_w8a8_moe_layer_for_cpu(
1625 + w13: torch.Tensor,
1626 + w2: torch.Tensor,
1627 + w13_scale: torch.Tensor,
1628 + w2_scale: torch.Tensor,
1629 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
1630 + """Prepack FP8 W8A8 MoE weights for CPU kernel."""
1631 + E = w13.size(0)
1632 +
1633 + packed_w13_list = []
1634 + packed_w13_scale_list = []
1635 +
1636 + for i in range(E):
1637 + w1_exp = w13[i] # [2N, K] FP8
1638 + ws_exp = w13_scale[i] # [2N, K/128] or [2N_blocks, K/128]
1639 +
1640 + # Align scale rows to weight rows if necessary (for block-128 checkpoints
1641 + # where scale may have fewer rows: [2N/128, K/128] → [2N, K/128])
1642 + if ws_exp.size(0) < w1_exp.size(0):
1643 + repeat_factor = w1_exp.size(0) // ws_exp.size(0)
1644 + ws_exp = torch.repeat_interleave(ws_exp, repeat_factor, dim=0)
1645 + ws_exp = ws_exp[: w1_exp.size(0), :].contiguous()
1646 +
1647 + pw, ps = torch.ops._C.float8_linear_prepack_cpu(w1_exp, ws_exp)
1648 + packed_w13_list.append(pw)
1649 + packed_w13_scale_list.append(ps)
1650 +
1651 + packed_w13 = torch.stack(packed_w13_list)
1652 + packed_w13_scale = torch.stack(packed_w13_scale_list)
1653 +
1654 + # w2 uses the W8A16 path (BF16 activation × FP8 weight), VNNI packed
1655 + packed_w2 = torch.ops._C.convert_weight_packed(w2)
1656 + return packed_w13, packed_w13_scale, packed_w2, w2_scale
1657 +
1658 +
1659 +class CPUExpertsFp8W8A8(mk.FusedMoEExpertsModular):
1660 + """CPU FP8 W8A8 block-quantized modular MoE experts."""
1661 +
1662 + def __init__(
1663 + self,
1664 + moe_config: FusedMoEConfig,
1665 + quant_config: FusedMoEQuantConfig,
1666 + ):
1667 + super().__init__(moe_config, quant_config)
1668 +
1669 + @property
1670 + def expects_unquantized_inputs(self) -> bool:
1671 + # Receives BF16 hidden states; quantization to FP8 is done inside apply()
1672 + return True
1673 +
1674 + @staticmethod
1675 + def activation_format() -> mk.FusedMoEActivationFormat:
1676 + return mk.FusedMoEActivationFormat.Standard
1677 +
1678 + @staticmethod
1679 + def _supports_current_device() -> bool:
1680 + # Requires native AMX-FP8 hardware.
1681 + return current_platform.is_cpu() and cpu_has_amx_fp8()
1682 +
1683 + @staticmethod
1684 + def _supports_no_act_and_mul() -> bool:
1685 + return False
1686 +
1687 + @staticmethod
1688 + def _supports_activation(activation: MoEActivation) -> bool:
1689 + return activation == MoEActivation.SILU
1690 +
1691 + @staticmethod
1692 + def _supports_parallel_config(
1693 + moe_parallel_config: FusedMoEParallelConfig,
1694 + ) -> bool:
1695 + return True
1696 +
1697 + @staticmethod
1698 + def _supports_quant_scheme(
1699 + weight_key: QuantKey | None,
1700 + activation_key: QuantKey | None,
1701 + ) -> bool:
1702 + # Block-quantized FP8.
1703 + return (weight_key, activation_key) == (
1704 + kFp8Static128BlockSym,
1705 + kFp8Dynamic128Sym,
1706 + )
1707 +
1708 + @staticmethod
1709 + def _supports_routing_method(
1710 + routing_method: RoutingMethodType,
1711 + weight_key: QuantKey | None,
1712 + activation_key: QuantKey | None,
1713 + ) -> bool:
1714 + return routing_method in [
1715 + RoutingMethodType.Default,
1716 + RoutingMethodType.Renormalize,
1717 + RoutingMethodType.RenormalizeNaive,
1718 + RoutingMethodType.DeepSeekV3,
1719 + RoutingMethodType.DeepseekV4,
1720 + ]
1721 +
1722 + @staticmethod
1723 + def _supports_router_logits_dtype(
1724 + router_logits_dtype: torch.dtype | None,
1725 + routing_method: RoutingMethodType,
1726 + ) -> bool:
1727 + return True
1728 +
1729 + def moe_problem_size(
1730 + self,
1731 + a1: torch.Tensor,
1732 + w1: torch.Tensor,
1733 + w2: torch.Tensor,
1734 + topk_ids: torch.Tensor,
1735 + ) -> tuple[int, int, int, int, int]:
1736 + # float8_linear_prepack_cpu blocks w1 into [E, Nc, Kc, block_k, block_n],
1737 + # not the (E, N, K) layout the base implementation assumes -- N isn't
1738 + # recoverable from that shape, so read it from moe_config instead.
1739 + E = w1.shape[0]
1740 + K = a1.size(-1)
1741 + N = (
1742 + self.moe_config.intermediate_size_per_partition
1743 + * self.moe_config.w13_num_shards
1744 + )
1745 + M = a1.size(0) if a1.dim() == 2 else a1.size(1)
1746 + topk = topk_ids.size(1)
1747 + return E, M, N, K, topk
1748 +
1749 + def workspace_shapes(
1750 + self,
1751 + M: int,
1752 + N: int,
1753 + K: int,
1754 + topk: int,
1755 + global_num_experts: int,
1756 + local_num_experts: int,
1757 + expert_tokens_meta: mk.ExpertTokensMetadata | None,
1758 + activation: MoEActivation,
1759 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
1760 + # fused_experts_cpu manages its own scratch space.
1761 + return (0,), (0,), (M, K)
1762 +
1763 + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
1764 + return TopKWeightAndReduceNoOP()
1765 +
1766 + def apply(
1767 + self,
1768 + output: torch.Tensor,
1769 + hidden_states: torch.Tensor,
1770 + w1: torch.Tensor,
1771 + w2: torch.Tensor,
1772 + topk_weights: torch.Tensor,
1773 + topk_ids: torch.Tensor,
1774 + activation: MoEActivation,
1775 + global_num_experts: int,
1776 + expert_map: torch.Tensor | None,
1777 + a1q_scale: torch.Tensor | None,
1778 + a2_scale: torch.Tensor | None,
1779 + workspace13: torch.Tensor,
1780 + workspace2: torch.Tensor,
1781 + expert_tokens_meta: mk.ExpertTokensMetadata | None,
1782 + apply_router_weight_on_input: bool,
1783 + ) -> None:
1784 + if apply_router_weight_on_input:
1785 + # fused_experts_cpu always applies topk_weights internally on
1786 + # combine; MoEPrepareAndFinalizeNoDPEPModular.prepare() would
1787 + # also pre-apply it to hidden_states, double-weighting the
1788 + # output. Not needed by any CPU FP8 model today.
1789 + raise NotImplementedError(
1790 + "CPUExpertsFp8W8A8 does not support apply_router_weight_on_input=True."
1791 + )
1792 +
1793 + block_shape = (
1794 + list(self.quant_config.block_shape)
1795 + if self.quant_config.block_shape
1796 + else (
1797 + [
1798 + self.quant_config._w1.shape.row,
1799 + self.quant_config._w1.shape.col,
1800 + ]
1801 + if self.quant_config._w1.shape is not None
1802 + else None
1803 + )
1804 + )
1805 +
1806 + # Quantize hidden_states (BF16) → FP8 with per-token scales
1807 + hidden_states_2d = hidden_states.reshape(-1, hidden_states.shape[-1])
1808 + x_fp8, x_scales = torch.ops._C.quantize_fp8e4m3_vec(
1809 + hidden_states_2d, True, None
1810 + )
1811 +
1812 + fused_experts_cpu(
1813 + output,
1814 + x_fp8,
1815 + w1,
1816 + w2,
1817 + topk_weights,
1818 + topk_ids,
1819 + CPUQuantMethod.FP8_W8A8,
1820 + self.w1_scale, # w1_scale
1821 + self.w2_scale, # w2_scale
1822 + None, # w1_zero
1823 + None, # w2_zero
1824 + block_shape, # block_size
1825 + None, # w1_bias
1826 + None, # w2_bias
1827 + None, # alpha
1828 + None, # limit
1829 + True, # is_vnni
1830 + x_scales, # a1_scale: per-token FP8 activation scales
1831 + )
vllm/model_executor/layers/fused_moe/oracle/fp8.py
+30 -2
@@ -57,7 +57,9 @@ class Fp8MoeBackend(Enum):
57 VLLM_CUTLASS = "VLLM_CUTLASS"
58 BATCHED_VLLM_CUTLASS = "BATCHED_VLLM_CUTLASS"
59 XPU = "XPU"
60 + # CPU FP8 W8A16 (BF16 activations); CPU_W8A8 needs native AMX-FP8.
61 CPU = "CPU"
62 + CPU_W8A8 = "CPU_W8A8"
63 HPC = "HPC"
64 # Dequantize-to-BF16 emulation for MXFP8 on devices without a native
65 # MXFP8 MoE kernel (e.g. ROCm). Weights pass through unchanged here.
@@ -92,6 +94,7 @@ def _get_priority_backends(
94 Fp8MoeBackend.BATCHED_VLLM_CUTLASS,
95 Fp8MoeBackend.BATCHED_TRITON,
96 Fp8MoeBackend.XPU,
97 + Fp8MoeBackend.CPU_W8A8,
98 Fp8MoeBackend.CPU,
99 Fp8MoeBackend.HPC,
100 ]
@@ -129,8 +132,10 @@ def _get_priority_backends(
132 _move_to_front(_AVAILABLE_BACKENDS, Fp8MoeBackend.XPU)
133
134 if current_platform.is_cpu():
132 - # CPU platform uses FP8 W8A16 fused MoE kernel.
135 + # W8A8 first: it falls through to the W8A16 backend whenever the
136 + # hardware (AMX-FP8) or the config isn't supported.
137 _move_to_front(_AVAILABLE_BACKENDS, Fp8MoeBackend.CPU)
138 + _move_to_front(_AVAILABLE_BACKENDS, Fp8MoeBackend.CPU_W8A8)
139
140 return _AVAILABLE_BACKENDS
141
@@ -238,6 +243,13 @@ def backend_to_kernel_cls(
243
244 return [CPUExpertsFp8]
245
246 + elif backend == Fp8MoeBackend.CPU_W8A8:
247 + from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
248 + CPUExpertsFp8W8A8,
249 + )
250 +
251 + return [CPUExpertsFp8W8A8]
252 +
253 elif backend == Fp8MoeBackend.HPC:
254 from vllm.model_executor.layers.fused_moe.hpc_moe import (
255 HPCExperts,
@@ -652,6 +664,14 @@ def convert_to_fp8_moe_kernel_format(
664 )
665
666 w13, w2 = prepare_fp8_moe_layer_for_cpu(w13, w2)
667 + elif fp8_backend == Fp8MoeBackend.CPU_W8A8:
668 + from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
669 + prepare_fp8_w8a8_moe_layer_for_cpu,
670 + )
671 +
672 + w13, w13_scale, w2, w2_scale = prepare_fp8_w8a8_moe_layer_for_cpu(
673 + w13, w2, w13_scale, w2_scale
674 + )
675 else:
676 if fp8_backend not in [
677 Fp8MoeBackend.TRITON,
@@ -697,7 +717,15 @@ def make_fp8_moe_quant_config(
717 In a future PR, we will have this function should be
718 a method of the modular kernel itself.
719 """
700 - # MARLIN and CPU are mixed precision W8A16 config.
720 + if fp8_backend == Fp8MoeBackend.CPU_W8A8:
721 + return fp8_w8a8_moe_quant_config(
722 + w1_scale=w1_scale,
723 + w2_scale=w2_scale,
724 + a1_scale=a1_scale,
725 + block_shape=block_shape,
726 + )
727 +
728 + # MARLIN and CPU (W8A16) are mixed precision W8A16 configs.
729 if fp8_backend == Fp8MoeBackend.MARLIN or fp8_backend == Fp8MoeBackend.CPU:
730 return fp8_w8a16_moe_quant_config(
731 w1_scale=w1_scale,
vllm/model_executor/layers/quantization/online/fp8.py
+4 -2
@@ -737,11 +737,13 @@ class Fp8PtpcOnlineMoEMethod(_Fp8OnlineMoEBase):
737 )
738 # Reject backends whose make_fp8_moe_quant_config branch silently
739 # drops per_act_token_quant / per_out_ch_quant or collapses scales:
740 - # MARLIN / CPU route through fp8_w8a16_moe_quant_config; FLASHINFER_*
741 - # fold scales into a per-tensor alpha (oracle/fp8.py).
740 + # MARLIN / CPU route through fp8_w8a16_moe_quant_config, CPU_W8A8
741 + # ignores both flags, and FLASHINFER_* fold scales into a per-tensor
742 + # alpha (oracle/fp8.py).
743 if self.fp8_backend in (
744 Fp8MoeBackend.MARLIN,
745 Fp8MoeBackend.CPU,
746 + Fp8MoeBackend.CPU_W8A8,
747 Fp8MoeBackend.FLASHINFER_CUTLASS,
748 Fp8MoeBackend.FLASHINFER_TRTLLM,
749 ):
vllm/v1/attention/backends/mla/amx_mla.py
+79 -8
@@ -61,6 +61,44 @@ _MIN_WORK_PER_SPLIT = 512
61 _SPLIT_OCCUPANCY_MULTIPLIER = 2
62
63
64 +def _block_quant_to_tensor_quant(
65 + weight: torch.Tensor, # FP8 [N, K] block-quantized weight
66 + weight_scale: torch.Tensor, # [ceil(N/bn), ceil(K/bk)] block scales
67 + block_size: list[int], # [block_n, block_k]
68 +) -> tuple[torch.Tensor, torch.Tensor]:
69 + """Re-quantize FP8 block-quant weight to FP8 per-tensor quant.
70 +
71 + Returns (fp8_weight, dequant_scale) where fp8_weight * dequant_scale
72 + approximates the original float weight. dequant_scale is a 0-dim
73 + scalar tensor as expected by bmm_cpu's ``scale`` argument.
74 +
75 + NOTE: block->tensor re-quantization may slightly affect accuracy
76 + (equivalent to SGLang's block_quant_to_tensor_quant for CPU).
77 + """
78 + block_n, block_k = block_size[0], block_size[1]
79 + n, k = weight.shape
80 +
81 + # Expand block scales [ceil(N/bn), ceil(K/bk)] -> [N, K]
82 + scale = weight_scale.to(torch.float32)
83 + if scale.shape[0] < n:
84 + scale = scale.repeat_interleave(block_n, dim=0)[:n, :]
85 + if scale.shape[1] < k:
86 + scale = scale.repeat_interleave(block_k, dim=1)[:, :k]
87 +
88 + # Dequantize to float32
89 + x_f32 = weight.to(torch.float32) * scale
90 +
91 + # Re-quantize per-tensor back to FP8
92 + fp8_max = torch.finfo(weight.dtype).max # 448.0 for e4m3fn
93 + amax = x_f32.abs().amax().clamp(min=1e-12)
94 + quant_scale = fp8_max / amax
95 + x_fp8 = (x_f32 * quant_scale).clamp(-fp8_max, fp8_max).to(weight.dtype)
96 +
97 + # dequant_scale = 1 / quant_scale (0-dim tensor for bmm_cpu)
98 + dequant_scale = (amax / fp8_max).to(torch.float32)
99 + return x_fp8, dequant_scale
100 +
101 +
102 def _compute_num_kv_splits(max_seq_len: int, num_threads: int) -> int:
103 """Mirrors TritonMLAImpl's _compute_num_kv_splits, using the CPU thread
104 count in place of SM count."""
@@ -208,6 +246,9 @@ class AMXMLAImpl(MLACommonImpl[MLACommonMetadata]):
246 )
247 self._w_uk_packed: torch.Tensor | None = None
248 self._w_uv_packed: torch.Tensor | None = None
249 + self._w_scale: torch.Tensor | None = (
250 + None # per-tensor dequant scale (FP8 W8A16)
251 + )
252
253 def process_weights_after_loading(self, act_dtype: torch.dtype) -> None:
254 # forward_mha (unlike forward_mqa) receives no `layer` argument, so it
@@ -216,12 +257,40 @@ class AMXMLAImpl(MLACommonImpl[MLACommonMetadata]):
257 # impl derives its own copies from self.kv_b_proj (already available
258 # via the constructor) and VNNI-packs them once, ahead of time, for
259 # use with bmm_cpu in forward_mha.
219 - kv_b_proj_weight = get_and_maybe_dequant_weights(
220 - self.kv_b_proj, out_dtype=act_dtype
221 - ).T
222 - kv_b_proj_weight = kv_b_proj_weight.view(
223 - self.kv_lora_rank, self.num_heads, self.qk_nope_head_dim + self.v_head_dim
224 - )
260 + #
261 + # kv_b_proj carries `_cpu_skip_gemm_dispatch`, so the CPU FP8 kernels
262 + # leave its weight in the checkpoint [N, K] layout. For block-quant
263 + # FP8 checkpoints, re-quantize per-tensor and keep the BMM in FP8
264 + # (W8A16); everything else goes through the BF16 path.
265 + weight = getattr(self.kv_b_proj, "weight", None)
266 + block_size = getattr(self.kv_b_proj, "weight_block_size", None)
267 + if (
268 + weight is not None
269 + and weight.dtype == torch.float8_e4m3fn
270 + and block_size is not None
271 + ):
272 + block_scale = getattr(self.kv_b_proj, "weight_scale_inv", None)
273 + if block_scale is None:
274 + block_scale = self.kv_b_proj.weight_scale
275 + fp8_weight, self._w_scale = _block_quant_to_tensor_quant(
276 + weight, block_scale, block_size
277 + )
278 + # [N, K] -> [K, N] = [kv_lora_rank, heads * (nope + v)]
279 + kv_b_proj_weight = fp8_weight.T.contiguous().view(
280 + self.kv_lora_rank,
281 + self.num_heads,
282 + self.qk_nope_head_dim + self.v_head_dim,
283 + )
284 + else:
285 + kv_b_proj_weight = get_and_maybe_dequant_weights(
286 + self.kv_b_proj, out_dtype=act_dtype
287 + ).T
288 + kv_b_proj_weight = kv_b_proj_weight.view(
289 + self.kv_lora_rank,
290 + self.num_heads,
291 + self.qk_nope_head_dim + self.v_head_dim,
292 + )
293 + self._w_scale = None
294 w_uk, w_uv = kv_b_proj_weight.split(
295 [self.qk_nope_head_dim, self.v_head_dim], dim=-1
296 )
@@ -348,7 +417,9 @@ class AMXMLAImpl(MLACommonImpl[MLACommonMetadata]):
417 ql_nope = torch.empty(
418 num_tokens, num_heads, self.kv_lora_rank, dtype=q.dtype, device=q.device
419 )
351 - ops.bmm_cpu(ql_nope.transpose(0, 1), q_nope_t, self._w_uk_packed, True, None)
420 + ops.bmm_cpu(
421 + ql_nope.transpose(0, 1), q_nope_t, self._w_uk_packed, True, self._w_scale
422 + )
423 mqa_q = torch.cat([ql_nope, q_pe], dim=-1)
424
425 kv_cache_flat = kv_c_and_k_pe_cache.view(-1, 1, self.head_size)
@@ -402,4 +473,4 @@ class AMXMLAImpl(MLACommonImpl[MLACommonMetadata]):
473 output_view = output.view(num_tokens, num_heads, self.v_head_dim).transpose(
474 0, 1
475 )
405 - ops.bmm_cpu(output_view, attn_out_t, self._w_uv_packed, True, None)
476 + ops.bmm_cpu(output_view, attn_out_t, self._w_uv_packed, True, self._w_scale)