| 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 | #ifdef CPU_CAPABILITY_AVX10_2 |
| 776 | __attribute__((target("avx10.2"))) inline __m128i cvtfp32_fp8e4m3_avx10_2(__m512& src) { |
| 777 | __m256i f16_vec = _mm512_cvt_roundps_ph(src, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC); |
| 778 | return _mm256_cvtph_hf8(_mm256_castsi256_ph(f16_vec)); |
| 779 | } |
| 780 | |
| 781 | // Carries the whole loop rather than just the conversion, as GCC cannot inline an avx10.2 |
| 782 | // function into a caller lacking that target, so branching per vector is costly. |
| 783 | __attribute__((target("avx10.2"))) void quantize_bf16_block_avx10_2( |
| 784 | const at::BFloat16* __restrict__ src, at::Float8_e4m3fn* __restrict__ dst, int64_t count, float scale_reciprocal) { |
| 785 | constexpr float quant_max = 448.0f; // torch.finfo(torch.float8_e4m3fn).max |
| 786 | const __m512 scale_recip_vec = _mm512_set1_ps(scale_reciprocal); |
| 787 | const __m512 quant_max_vec = _mm512_set1_ps(quant_max); |
| 788 | const __m512 neg_quant_max_vec = _mm512_set1_ps(-quant_max); |
| 789 | |
| 790 | for (int64_t i = 0; i <= count - 32; i += 32) { |
| 791 | __m256i src_vec1 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&src[i])); |
| 792 | __m256i src_vec2 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&src[i + 16])); |
| 793 | |
| 794 | __m512 fp32_vec1 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(src_vec1), 16)); |
| 795 | __m512 fp32_vec2 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(src_vec2), 16)); |
| 796 | |
| 797 | __m512 scaled_vec1 = _mm512_mul_ps(fp32_vec1, scale_recip_vec); |
| 798 | __m512 clamped_vec1 = _mm512_min_ps(_mm512_max_ps(scaled_vec1, neg_quant_max_vec), quant_max_vec); |
| 799 | |
| 800 | __m512 scaled_vec2 = _mm512_mul_ps(fp32_vec2, scale_recip_vec); |
| 801 | __m512 clamped_vec2 = _mm512_min_ps(_mm512_max_ps(scaled_vec2, neg_quant_max_vec), quant_max_vec); |
| 802 | |
| 803 | __m128i fp8_vec1 = cvtfp32_fp8e4m3_avx10_2(clamped_vec1); |
| 804 | __m128i fp8_vec2 = cvtfp32_fp8e4m3_avx10_2(clamped_vec2); |
| 805 | _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + i), fp8_vec1); |
| 806 | _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + i + 16), fp8_vec2); |
| 807 | } |
| 808 | } |
| 809 | #endif // CPU_CAPABILITY_AVX10_2 |
| 810 | |
| 811 | std::tuple<at::Tensor, at::Tensor> _quantize_fp8e4m3_bf16_per_tensor_no_scale(const at::Tensor& t) { |
| 812 | constexpr float quant_max = 448.0f; // torch.finfo(torch.float8_e4m3fn).max |
| 813 | constexpr float eps = std::numeric_limits<float>::epsilon(); |
| 814 | |
| 815 | // Input validation and preparation |
| 816 | assert(t.scalar_type() == at::ScalarType::BFloat16); |
| 817 | auto t_bf16 = t.contiguous(); |
| 818 | int64_t num_channels = t_bf16.size(0); |
| 819 | int64_t elements_per_channel = t_bf16.numel() / num_channels; |
| 820 | assert(elements_per_channel % 32 == 0); // do not consider tile currently |
| 821 | at::Tensor quant_t = at::empty(t_bf16.sizes(), t_bf16.options().dtype(at::kFloat8_e4m3fn)); |
| 822 | |
| 823 | // Allocate output tensors |
| 824 | at::Tensor scale_tensor = at::empty({num_channels}, t_bf16.options().dtype(at::ScalarType::Float)); |
| 825 | float* scale_data = scale_tensor.data_ptr<float>(); |
| 826 | |
| 827 | // Unified processing: compute max, apply scale, and quantize in single pass |
| 828 | at::parallel_for(0, num_channels, 1, [&](int64_t start, int64_t end) { |
| 829 | for (int64_t c = start; c < end; ++c) { |
| 830 | const at::BFloat16* channel_src = t_bf16.data_ptr<at::BFloat16>() + c * elements_per_channel; |
| 831 | at::Float8_e4m3fn* quant_dst = quant_t.data_ptr<at::Float8_e4m3fn>() + c * elements_per_channel; |
| 832 | |
| 833 | // Step 1: Compute channel-wise max using AVX512 |
| 834 | float channel_max = 0.0f; |
| 835 | int64_t i = 0; |
| 836 | |
| 837 | // Process 32 elements at a time for max computation |
| 838 | for (; i <= elements_per_channel - 32; i += 32) { |
| 839 | // Load 32 BF16 values (2x 256-bit vectors) |
| 840 | __m256i src_vec1 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&channel_src[i])); |
| 841 | __m256i src_vec2 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&channel_src[i + 16])); |
| 842 | |
| 843 | // Convert BF16 to FP32 |
| 844 | __m512i fp32_int1 = _mm512_cvtepu16_epi32(src_vec1); |
| 845 | __m512 fp32_vec1 = _mm512_castsi512_ps(_mm512_slli_epi32(fp32_int1, 16)); |
| 846 | |
| 847 | __m512i fp32_int2 = _mm512_cvtepu16_epi32(src_vec2); |
| 848 | __m512 fp32_vec2 = _mm512_castsi512_ps(_mm512_slli_epi32(fp32_int2, 16)); |
| 849 | |
| 850 | // Compute absolute values |
| 851 | __m512 abs_vec1 = _mm512_abs_ps(fp32_vec1); |
| 852 | __m512 abs_vec2 = _mm512_abs_ps(fp32_vec2); |
| 853 | |
| 854 | // Find max in each vector |
| 855 | float max1 = _mm512_reduce_max_ps(abs_vec1); |
| 856 | float max2 = _mm512_reduce_max_ps(abs_vec2); |
| 857 | channel_max = std::max(channel_max, std::max(max1, max2)); |
| 858 | } |
| 859 | |
| 860 | // Step 2: Apply quant_max and EPS to compute scale |
| 861 | float scale_val = std::max(channel_max / quant_max, eps); |
| 862 | float scale_reciprocal = 1.0f / scale_val; |
| 863 | scale_data[c] = scale_val; |
| 864 | |
| 865 | // Step 3: Scale and clamp using AVX512 (reuse the same loop structure) |
| 866 | #ifdef CPU_CAPABILITY_AVX10_2 |
| 867 | if (avx10_2_available()) { |
| 868 | quantize_bf16_block_avx10_2(channel_src, quant_dst, elements_per_channel, scale_reciprocal); |
| 869 | continue; |
| 870 | } |
| 871 | #endif |
| 872 | i = 0; |
| 873 | const __m512 scale_recip_vec = _mm512_set1_ps(scale_reciprocal); |
| 874 | const __m512 quant_max_vec = _mm512_set1_ps(quant_max); |
| 875 | const __m512 neg_quant_max_vec = _mm512_set1_ps(-quant_max); |
| 876 | |
| 877 | for (; i <= elements_per_channel - 32; i += 32) { |
| 878 | // Load 32 BF16 values |
| 879 | __m256i src_vec1 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&channel_src[i])); |
| 880 | __m256i src_vec2 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&channel_src[i + 16])); |
| 881 | |
| 882 | // Convert BF16 to FP32 |
| 883 | __m512 fp32_vec1 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(src_vec1), 16)); |
| 884 | __m512 fp32_vec2 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(src_vec2), 16)); |
| 885 | |
| 886 | // Scale and clamp |
| 887 | __m512 scaled_vec1 = _mm512_mul_ps(fp32_vec1, scale_recip_vec); |
| 888 | __m512 clamped_vec1 = _mm512_min_ps(_mm512_max_ps(scaled_vec1, neg_quant_max_vec), quant_max_vec); |
| 889 | |
| 890 | __m512 scaled_vec2 = _mm512_mul_ps(fp32_vec2, scale_recip_vec); |
| 891 | __m512 clamped_vec2 = _mm512_min_ps(_mm512_max_ps(scaled_vec2, neg_quant_max_vec), quant_max_vec); |
| 892 | |
| 893 | __m128i fp8_vec1 = cvtfp32_fp8e4m3(clamped_vec1); |
| 894 | __m128i fp8_vec2 = cvtfp32_fp8e4m3(clamped_vec2); |
| 895 | _mm_storeu_si128(reinterpret_cast<__m128i*>(quant_dst + i), fp8_vec1); |
| 896 | _mm_storeu_si128(reinterpret_cast<__m128i*>(quant_dst + i + 16), fp8_vec2); |
| 897 | } |
| 898 | } |
| 899 | }); |
| 900 | |
| 901 | return std::make_tuple(quant_t, scale_tensor); |
| 902 | } |
| 903 | |
| 904 | std::tuple<at::Tensor, at::Tensor> |
| 905 | _quantize_fp8e4m3_bf16_per_tensor_with_scale(const at::Tensor& t, at::Tensor& scale_tensor) { |
| 906 | constexpr float quant_max = 448.0f; // torch.finfo(torch.float8_e4m3fn).max |
| 907 | constexpr float eps = std::numeric_limits<float>::epsilon(); |
| 908 | |
| 909 | // Input validation |
| 910 | assert(t.scalar_type() == at::ScalarType::BFloat16); |
| 911 | |
| 912 | auto t_bf16 = t.contiguous(); |
| 913 | auto scale_tensor_contig = scale_tensor.contiguous().to(at::ScalarType::Float); |
| 914 | |
| 915 | // Get scale value (handle scalar or tensor) |
| 916 | float scale_val; |
| 917 | if (scale_tensor_contig.numel() == 1) { |
| 918 | scale_val = scale_tensor_contig.item<float>(); |
| 919 | } else { |
| 920 | // Take max if scale is a tensor (though should be scalar for global quantization) |
| 921 | scale_val = scale_tensor_contig.max().item<float>(); |
| 922 | } |
| 923 | |
| 924 | // Apply EPS to scale |
| 925 | scale_val = std::max(scale_val, eps); |
| 926 | float scale_reciprocal = 1.0f / scale_val; |
| 927 | |
| 928 | at::Tensor scale_output = at::tensor({scale_val}, scale_tensor_contig.options()); |
| 929 | |
| 930 | // Apply scale and clamp using AVX512 |
| 931 | const at::BFloat16* src_data = t_bf16.data_ptr<at::BFloat16>(); |
| 932 | at::Tensor quant_t = at::empty(t_bf16.sizes(), t_bf16.options().dtype(at::kFloat8_e4m3fn)); |
| 933 | int64_t total_elements = t_bf16.numel(); |
| 934 | const __m512 scale_recip_vec = _mm512_set1_ps(scale_reciprocal); |
| 935 | const __m512 quant_max_vec = _mm512_set1_ps(quant_max); |
| 936 | const __m512 neg_quant_max_vec = _mm512_set1_ps(-quant_max); |
| 937 | int64_t num_channels = t_bf16.size(0); |
| 938 | int64_t elements_per_channel = total_elements / num_channels; |
| 939 | assert(elements_per_channel % 32 == 0); // do not consider tile currently |
| 940 | // Process in parallel blocks |
| 941 | at::parallel_for(0, num_channels, 1, [&](int64_t start, int64_t end) { |
| 942 | for (int64_t c = start; c < end; ++c) { |
| 943 | const at::BFloat16* src_data = t_bf16.data_ptr<at::BFloat16>() + c * elements_per_channel; |
| 944 | at::Float8_e4m3fn* quant_t_data = quant_t.data_ptr<at::Float8_e4m3fn>() + c * elements_per_channel; |
| 945 | |
| 946 | #ifdef CPU_CAPABILITY_AVX10_2 |
| 947 | if (avx10_2_available()) { |
| 948 | quantize_bf16_block_avx10_2(src_data, quant_t_data, elements_per_channel, scale_reciprocal); |
| 949 | continue; |
| 950 | } |
| 951 | #endif |
| 952 | |
| 953 | int64_t i = 0; |
| 954 | // Process 32 elements at a time using AVX512 |
| 955 | for (; i <= elements_per_channel - 32; i += 32) { |
| 956 | // Load 32 BF16 values (2x 256-bit vectors) |
| 957 | __m256i src_vec1 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&src_data[i])); |
| 958 | __m256i src_vec2 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&src_data[i + 16])); |
| 959 | |
| 960 | // Convert BF16 to FP32 |
| 961 | __m512 fp32_vec1 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(src_vec1), 16)); |
| 962 | __m512 fp32_vec2 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(src_vec2), 16)); |
| 963 | |
| 964 | // Scale and clamp |
| 965 | __m512 scaled_vec1 = _mm512_mul_ps(fp32_vec1, scale_recip_vec); |
| 966 | __m512 clamped_vec1 = _mm512_min_ps(_mm512_max_ps(scaled_vec1, neg_quant_max_vec), quant_max_vec); |
| 967 | |
| 968 | __m512 scaled_vec2 = _mm512_mul_ps(fp32_vec2, scale_recip_vec); |
| 969 | __m512 clamped_vec2 = _mm512_min_ps(_mm512_max_ps(scaled_vec2, neg_quant_max_vec), quant_max_vec); |
| 970 | |
| 971 | __m128i fp8_vec1 = cvtfp32_fp8e4m3(clamped_vec1); |
| 972 | __m128i fp8_vec2 = cvtfp32_fp8e4m3(clamped_vec2); |
| 973 | _mm_storeu_si128(reinterpret_cast<__m128i*>(quant_t_data + i), fp8_vec1); |
| 974 | _mm_storeu_si128(reinterpret_cast<__m128i*>(quant_t_data + i + 16), fp8_vec2); |
| 975 | } |
| 976 | } |
| 977 | }); |
| 978 | |
| 979 | return std::make_tuple(quant_t, scale_output); |
| 980 | } |
| 981 | |
| 982 | std::tuple<at::Tensor, at::Tensor> |
| 983 | _quantize_fp8e4m3_vec(const at::Tensor& t, bool channelwise, c10::optional<at::Tensor> scale_opt) { |
| 984 | if (channelwise) { |
| 985 | return _quantize_fp8e4m3_bf16_per_tensor_no_scale(t); |
| 986 | } else { |
| 987 | assert(scale_opt.has_value()); |
| 988 | return _quantize_fp8e4m3_bf16_per_tensor_with_scale(t, scale_opt.value()); |
| 989 | } |
| 990 | } |
| 991 | |
| 992 | // Public wrapper for torch op registration |
| 993 | std::tuple<at::Tensor, at::Tensor> |
| 994 | quantize_fp8e4m3_vec(const at::Tensor& t, bool channelwise, c10::optional<at::Tensor> scale_opt) { |
| 995 | return _quantize_fp8e4m3_vec(t, channelwise, scale_opt); |
| 996 | } |
| 997 | |
| 998 | at::Tensor fp8_scaled_mm_with_quant( |
| 999 | const at::Tensor& act, |
| 1000 | const std::optional<at::Tensor>& act_scales, |
| 1001 | bool channelwise, |
| 1002 | const at::Tensor& weight, |
| 1003 | const at::Tensor& weight_scales, |
| 1004 | const std::optional<at::Tensor>& bias, |
| 1005 | at::ScalarType output_dtype) { |
| 1006 | // TODO: refine here to use tensor ptr and in below dispatch api. |
| 1007 | std::tuple<at::Tensor, at::Tensor> quant_act = act.scalar_type() == at::ScalarType::BFloat16 |
| 1008 | ? _quantize_fp8e4m3_vec(act, channelwise, act_scales) |
| 1009 | : _quantize_fp8e4m3(act, channelwise, act_scales); |
| 1010 | auto input = std::get<0>(quant_act); |
| 1011 | auto input_scales = std::get<1>(quant_act); |
| 1012 | int64_t N = weight.dim() == 4 ? weight.size(0) * weight.size(-1) : weight.size(0); |
| 1013 | int act_quant_mode = input_scales.numel() == 1 ? PER_TENSOR |
| 1014 | : input_scales.numel() == input.numel() / input.size(-1) ? PER_ROW |
| 1015 | : PER_GROUP; |
| 1016 | int wei_quant_mode = weight_scales.numel() == 1 ? PER_TENSOR : weight_scales.numel() == N ? PER_ROW : PER_GROUP; |
| 1017 | // Case to fall back |
| 1018 | if (weight.dim() == 2) { |
| 1019 | TORCH_CHECK( |
| 1020 | act_quant_mode != PER_GROUP && wei_quant_mode != PER_GROUP, |
| 1021 | "FP8 linear: Per-group quantization is not supported in the fallback path"); |
| 1022 | auto y_fp32 = at::linear(input.to(at::kFloat).mul_(input_scales), weight.to(at::kFloat).mul_(weight_scales), bias); |
| 1023 | return y_fp32.to(output_dtype); |
| 1024 | } |
| 1025 | |
| 1026 | static bool cpublas_can_pack = cpublas_could_pack(); |
| 1027 | auto out_sizes = input.sizes().vec(); |
| 1028 | out_sizes.back() = N; |
| 1029 | auto output = at::empty(out_sizes, input.options().dtype(output_dtype)); |
| 1030 | |
| 1031 | #define AT_DISPATCH_FP8_LINEAR_KERNEL(OUT_DTYPE, CAN_PACK, A_QUANT_MODE, B_QUANT_MODE, ...) \ |
| 1032 | AT_DISPATCH_BOOL_NO_RETURN( \ |
| 1033 | CAN_PACK, \ |
| 1034 | "cpublas_can_pack", \ |
| 1035 | can_pack, \ |
| 1036 | AT_DISPATCH_QUANT_MODE_NO_RETURN( \ |
| 1037 | A_QUANT_MODE, \ |
| 1038 | "act_quant_mode", \ |
| 1039 | a_quant_mode, \ |
| 1040 | AT_DISPATCH_QUANT_MODE_NO_RETURN( \ |
| 1041 | B_QUANT_MODE, \ |
| 1042 | "wei_quant_mode", \ |
| 1043 | b_quant_mode, \ |
| 1044 | AT_DISPATCH_OUT_TYPES(OUT_DTYPE, "out_dtype", __VA_ARGS__)))) |
| 1045 | |
| 1046 | AT_DISPATCH_FP8_LINEAR_KERNEL(output_dtype, cpublas_can_pack, act_quant_mode, wei_quant_mode, [&]() { |
| 1047 | _float8_linear_impl<out_t, can_pack, a_quant_mode, b_quant_mode>( |
| 1048 | input, input_scales, weight, weight_scales, bias, output); |
| 1049 | }); |
| 1050 | return output; |
| 1051 | } |
| 1052 | |
| 1053 | at::Tensor float8_linear_impl( |
| 1054 | const at::Tensor& input, |
| 1055 | const at::Tensor& input_scales, |
| 1056 | const at::Tensor& weight, |
| 1057 | const at::Tensor& weight_scales, |
| 1058 | const std::optional<at::Tensor>& bias, |
| 1059 | at::ScalarType output_dtype) { |
| 1060 | int64_t N = weight.dim() == 4 ? weight.size(0) * weight.size(-1) : weight.size(0); |
| 1061 | int act_quant_mode = input_scales.numel() == 1 ? PER_TENSOR |
| 1062 | : input_scales.numel() == input.numel() / input.size(-1) ? PER_ROW |
| 1063 | : PER_GROUP; |
| 1064 | int wei_quant_mode = weight_scales.numel() == 1 ? PER_TENSOR : weight_scales.numel() == N ? PER_ROW : PER_GROUP; |
| 1065 | // Case to fall back |
| 1066 | if (weight.dim() == 2) { |
| 1067 | TORCH_CHECK( |
| 1068 | act_quant_mode != PER_GROUP && wei_quant_mode != PER_GROUP, |
| 1069 | "FP8 linear: Per-group quantization is not supported in the fallback path"); |
| 1070 | auto y_fp32 = at::linear(input.to(at::kFloat).mul_(input_scales), weight.to(at::kFloat).mul_(weight_scales), bias); |
| 1071 | return y_fp32.to(output_dtype); |
| 1072 | } |
| 1073 | |
| 1074 | static bool cpublas_can_pack = cpublas_could_pack(); |
| 1075 | auto out_sizes = input.sizes().vec(); |
| 1076 | out_sizes.back() = N; |
| 1077 | auto output = at::empty(out_sizes, input.options().dtype(output_dtype)); |
| 1078 | |
| 1079 | #define AT_DISPATCH_FP8_LINEAR_KERNEL(OUT_DTYPE, CAN_PACK, A_QUANT_MODE, B_QUANT_MODE, ...) \ |
| 1080 | AT_DISPATCH_BOOL_NO_RETURN( \ |
| 1081 | CAN_PACK, \ |
| 1082 | "cpublas_can_pack", \ |
| 1083 | can_pack, \ |
| 1084 | AT_DISPATCH_QUANT_MODE_NO_RETURN( \ |
| 1085 | A_QUANT_MODE, \ |
| 1086 | "act_quant_mode", \ |
| 1087 | a_quant_mode, \ |
| 1088 | AT_DISPATCH_QUANT_MODE_NO_RETURN( \ |
| 1089 | B_QUANT_MODE, \ |
| 1090 | "wei_quant_mode", \ |
| 1091 | b_quant_mode, \ |
| 1092 | AT_DISPATCH_OUT_TYPES(OUT_DTYPE, "out_dtype", __VA_ARGS__)))) |
| 1093 | |
| 1094 | AT_DISPATCH_FP8_LINEAR_KERNEL(output_dtype, cpublas_can_pack, act_quant_mode, wei_quant_mode, [&]() { |
| 1095 | _float8_linear_impl<out_t, can_pack, a_quant_mode, b_quant_mode>( |
| 1096 | input, input_scales, weight, weight_scales, bias, output); |
| 1097 | }); |
| 1098 | return output; |
| 1099 | } |