| 1 | #include "cpu_types.hpp" |
| 2 | #include "dnnl_helper.h" |
| 3 | |
| 4 | namespace { |
| 5 | template <typename scalar_t> |
| 6 | struct KernelVecType { |
| 7 | using load_vec_type = void; |
| 8 | using cvt_vec_type = void; |
| 9 | }; |
| 10 | |
| 11 | template <> |
| 12 | struct KernelVecType<float> { |
| 13 | using load_vec_type = vec_op::FP32Vec16; |
| 14 | using cvt_vec_type = vec_op::FP32Vec16; |
| 15 | }; |
| 16 | |
| 17 | template <> |
| 18 | struct KernelVecType<c10::BFloat16> { |
| 19 | using load_vec_type = vec_op::BF16Vec16; |
| 20 | using cvt_vec_type = vec_op::FP32Vec16; |
| 21 | }; |
| 22 | |
| 23 | template <> |
| 24 | struct KernelVecType<c10::Half> { |
| 25 | #if defined(__powerpc64__) || defined(__s390x__) |
| 26 | // Power architecture-specific vector type |
| 27 | using load_vec_type = vec_op::FP32Vec16; |
| 28 | #else |
| 29 | // Fallback for other architectures |
| 30 | using load_vec_type = vec_op::FP16Vec16; |
| 31 | #endif |
| 32 | using cvt_vec_type = vec_op::FP32Vec16; |
| 33 | }; |
| 34 | |
| 35 | template <bool AZP, typename scalar_t> |
| 36 | void static_scaled_int8_quant_impl(const scalar_t* input, int8_t* output, |
| 37 | const float* scale, const int32_t* azp, |
| 38 | const int64_t num_tokens, |
| 39 | const int64_t input_stride, |
| 40 | const int64_t hidden_size) { |
| 41 | using load_vec_t = typename KernelVecType<scalar_t>::load_vec_type; |
| 42 | using cvt_vec_t = typename KernelVecType<scalar_t>::cvt_vec_type; |
| 43 | constexpr int64_t vec_elem_num = load_vec_t::VEC_ELEM_NUM; |
| 44 | |
| 45 | constexpr float i8_min = |
| 46 | static_cast<float>(std::numeric_limits<int8_t>::min()); |
| 47 | constexpr float i8_max = |
| 48 | static_cast<float>(std::numeric_limits<int8_t>::max()); |
| 49 | const cvt_vec_t inv_scale(1.0 / *scale); |
| 50 | const cvt_vec_t i8_min_vec(i8_min); |
| 51 | const cvt_vec_t i8_max_vec(i8_max); |
| 52 | |
| 53 | cvt_vec_t zp_vec; |
| 54 | if constexpr (AZP) { |
| 55 | zp_vec = cvt_vec_t(static_cast<float>(*azp)); |
| 56 | } |
| 57 | |
| 58 | #pragma omp parallel for |
| 59 | for (int64_t i = 0; i < num_tokens; ++i) { |
| 60 | int64_t j = 0; |
| 61 | const scalar_t* input_ptr = input + i * input_stride; |
| 62 | int8_t* output_ptr = output + i * hidden_size; |
| 63 | for (; j < hidden_size - vec_elem_num; j += vec_elem_num) { |
| 64 | load_vec_t elems(input_ptr + j); |
| 65 | cvt_vec_t elems_fp32(elems); |
| 66 | elems_fp32 = elems_fp32 * inv_scale; |
| 67 | |
| 68 | if constexpr (AZP) { |
| 69 | elems_fp32 = elems_fp32 + zp_vec; |
| 70 | } |
| 71 | |
| 72 | elems_fp32 = elems_fp32.clamp(i8_min_vec, i8_max_vec); |
| 73 | vec_op::INT8Vec16 elems_int8(elems_fp32); |
| 74 | elems_int8.save(output_ptr + j); |
| 75 | } |
| 76 | |
| 77 | load_vec_t elems(input_ptr + j); |
| 78 | cvt_vec_t elems_fp32(elems); |
| 79 | elems_fp32 = elems_fp32 * inv_scale; |
| 80 | |
| 81 | if constexpr (AZP) { |
| 82 | elems_fp32 = elems_fp32 + zp_vec; |
| 83 | } |
| 84 | |
| 85 | elems_fp32 = elems_fp32.clamp(i8_min_vec, i8_max_vec); |
| 86 | vec_op::INT8Vec16 elems_int8(elems_fp32); |
| 87 | elems_int8.save(output_ptr + j, hidden_size - j); |
| 88 | } |
| 89 | } |
| 90 | |
| 91 | template <bool AZP, typename scalar_t> |
| 92 | void dynamic_scaled_int8_quant_impl(const scalar_t* input, int8_t* output, |
| 93 | float* scale, int32_t* azp, |
| 94 | const int64_t num_tokens, |
| 95 | const int64_t input_stride, |
| 96 | const int64_t hidden_size) { |
| 97 | using load_vec_t = typename KernelVecType<scalar_t>::load_vec_type; |
| 98 | using cvt_vec_t = typename KernelVecType<scalar_t>::cvt_vec_type; |
| 99 | constexpr int vec_elem_num = load_vec_t::VEC_ELEM_NUM; |
| 100 | |
| 101 | constexpr float i8_min = |
| 102 | static_cast<float>(std::numeric_limits<int8_t>::min()); |
| 103 | constexpr float i8_max = |
| 104 | static_cast<float>(std::numeric_limits<int8_t>::max()); |
| 105 | const cvt_vec_t i8_min_vec(i8_min); |
| 106 | const cvt_vec_t i8_max_vec(i8_max); |
| 107 | |
| 108 | #pragma omp parallel for |
| 109 | for (int64_t i = 0; i < num_tokens; ++i) { |
| 110 | cvt_vec_t max_value(std::numeric_limits<float>::lowest()); |
| 111 | cvt_vec_t min_value(std::numeric_limits<float>::max()); |
| 112 | { |
| 113 | int64_t j = 0; |
| 114 | const scalar_t* input_ptr = input + i * input_stride; |
| 115 | for (; j < hidden_size - vec_elem_num; j += vec_elem_num) { |
| 116 | load_vec_t elems(input_ptr + j); |
| 117 | cvt_vec_t elems_fp32(elems); |
| 118 | if constexpr (AZP) { |
| 119 | max_value = max_value.max(elems_fp32); |
| 120 | min_value = min_value.min(elems_fp32); |
| 121 | } else { |
| 122 | max_value = max_value.max(elems_fp32.abs()); |
| 123 | } |
| 124 | } |
| 125 | |
| 126 | load_vec_t elems(input_ptr + j); |
| 127 | cvt_vec_t elems_fp32(elems); |
| 128 | |
| 129 | if (j + vec_elem_num == hidden_size) { |
| 130 | if constexpr (AZP) { |
| 131 | max_value = max_value.max(elems_fp32); |
| 132 | min_value = min_value.min(elems_fp32); |
| 133 | } else { |
| 134 | max_value = max_value.max(elems_fp32.abs()); |
| 135 | } |
| 136 | } else { |
| 137 | if constexpr (AZP) { |
| 138 | max_value = max_value.max(elems_fp32, hidden_size - j); |
| 139 | min_value = min_value.min(elems_fp32, hidden_size - j); |
| 140 | } else { |
| 141 | max_value = max_value.max(elems_fp32.abs(), hidden_size - j); |
| 142 | } |
| 143 | } |
| 144 | } |
| 145 | |
| 146 | float scale_val; |
| 147 | float azp_val = 0.0f; |
| 148 | if constexpr (AZP) { |
| 149 | float max_scalar = max_value.reduce_max(); |
| 150 | float min_scalar = min_value.reduce_min(); |
| 151 | scale_val = (max_scalar - min_scalar) / 255.0f; |
| 152 | azp_val = std::nearbyint(-128.0f - min_scalar / scale_val); |
| 153 | azp[i] = azp_val; |
| 154 | scale[i] = scale_val; |
| 155 | } else { |
| 156 | scale_val = max_value.reduce_max() / 127.0f; |
| 157 | scale[i] = scale_val; |
| 158 | } |
| 159 | |
| 160 | const cvt_vec_t inv_scale(1.0 / scale_val); |
| 161 | const cvt_vec_t azp_vec(azp_val); |
| 162 | |
| 163 | { |
| 164 | int64_t j = 0; |
| 165 | const scalar_t* input_ptr = input + i * input_stride; |
| 166 | int8_t* output_ptr = output + i * hidden_size; |
| 167 | for (; j < hidden_size - vec_elem_num; j += vec_elem_num) { |
| 168 | load_vec_t elems(input_ptr + j); |
| 169 | cvt_vec_t elems_fp32(elems); |
| 170 | elems_fp32 = (elems_fp32 * inv_scale); |
| 171 | |
| 172 | if constexpr (AZP) { |
| 173 | elems_fp32 = elems_fp32 + azp_vec; |
| 174 | } |
| 175 | elems_fp32 = elems_fp32.clamp(i8_min_vec, i8_max_vec); |
| 176 | vec_op::INT8Vec16 elems_int8(elems_fp32); |
| 177 | elems_int8.save(output_ptr + j); |
| 178 | } |
| 179 | |
| 180 | load_vec_t elems(input_ptr + j); |
| 181 | cvt_vec_t elems_fp32(elems); |
| 182 | elems_fp32 = (elems_fp32 * inv_scale); |
| 183 | |
| 184 | if constexpr (AZP) { |
| 185 | elems_fp32 = elems_fp32 + azp_vec; |
| 186 | } |
| 187 | elems_fp32 = elems_fp32.clamp(i8_min_vec, i8_max_vec); |
| 188 | vec_op::INT8Vec16 elems_int8(elems_fp32); |
| 189 | elems_int8.save(output_ptr + j, hidden_size - j); |
| 190 | } |
| 191 | } |
| 192 | } |
| 193 | |
| 194 | template <bool AZP, bool Bias, typename scalar_t> |
| 195 | void dynamic_quant_epilogue(const float* input, scalar_t* output, |
| 196 | const float* a_scale, const int32_t* azp, |
| 197 | const float* azp_adj, const scalar_t* bias, |
| 198 | const int64_t num_tokens, |
| 199 | const int64_t hidden_size) { |
| 200 | CPU_KERNEL_GUARD_IN(dynamic_quant_epilogue) |
| 201 | using load_vec_t = typename KernelVecType<scalar_t>::load_vec_type; |
| 202 | using cvt_vec_t = typename KernelVecType<scalar_t>::cvt_vec_type; |
| 203 | constexpr int vec_elem_num = load_vec_t::VEC_ELEM_NUM; |
| 204 | |
| 205 | const int64_t thread_num = cpu_utils::get_max_threads(); |
| 206 | if (num_tokens > thread_num) { |
| 207 | #pragma omp parallel for |
| 208 | for (int64_t i = 0; i < num_tokens; ++i) { |
| 209 | const float* input_ptr = input + i * hidden_size; |
| 210 | scalar_t* output_ptr = output + i * hidden_size; |
| 211 | int64_t j = 0; |
| 212 | cvt_vec_t token_scale_vec(a_scale[i]); |
| 213 | cvt_vec_t token_zp_scale_vec; |
| 214 | if constexpr (AZP) { |
| 215 | float zp_scale_val = a_scale[i] * static_cast<float>(azp[i]); |
| 216 | token_zp_scale_vec = cvt_vec_t(zp_scale_val); |
| 217 | } |
| 218 | for (; j < hidden_size - vec_elem_num; j += vec_elem_num) { |
| 219 | cvt_vec_t elems_fp32(input_ptr + j); |
| 220 | elems_fp32 = elems_fp32 * token_scale_vec; |
| 221 | if constexpr (AZP) { |
| 222 | cvt_vec_t azp_adj_fp32(azp_adj + j); |
| 223 | elems_fp32 = elems_fp32 - azp_adj_fp32 * token_zp_scale_vec; |
| 224 | } |
| 225 | if constexpr (Bias) { |
| 226 | load_vec_t bias_vec(bias + j); |
| 227 | cvt_vec_t bias_vec_fp32(bias_vec); |
| 228 | elems_fp32 = elems_fp32 + bias_vec_fp32; |
| 229 | } |
| 230 | load_vec_t elems_out(elems_fp32); |
| 231 | elems_out.save(output_ptr + j); |
| 232 | } |
| 233 | cvt_vec_t elems_fp32(input_ptr + j); |
| 234 | elems_fp32 = elems_fp32 * token_scale_vec; |
| 235 | if constexpr (AZP) { |
| 236 | cvt_vec_t azp_adj_fp32(azp_adj + j); |
| 237 | elems_fp32 = elems_fp32 - azp_adj_fp32 * token_zp_scale_vec; |
| 238 | } |
| 239 | if constexpr (Bias) { |
| 240 | load_vec_t bias_vec(bias + j); |
| 241 | cvt_vec_t bias_vec_fp32(bias_vec); |
| 242 | elems_fp32 = elems_fp32 + bias_vec_fp32; |
| 243 | } |
| 244 | load_vec_t elems_out(elems_fp32); |
| 245 | elems_out.save(output_ptr + j, hidden_size - j); |
| 246 | } |
| 247 | } else { |
| 248 | const int64_t vec_iteration = |
| 249 | (hidden_size + vec_elem_num - 1) / vec_elem_num; |
| 250 | const int64_t vec_iteration_per_thread = |
| 251 | (vec_iteration + thread_num - 1) / thread_num; |
| 252 | const int64_t elem_num_per_thread = vec_iteration_per_thread * vec_elem_num; |
| 253 | #pragma omp parallel for schedule(static, 1) |
| 254 | for (int64_t i = 0; i < thread_num; ++i) { |
| 255 | const int64_t start = elem_num_per_thread * i; |
| 256 | const int64_t end = std::min(hidden_size, elem_num_per_thread + start); |
| 257 | for (int64_t j = 0; j < num_tokens; ++j) { |
| 258 | cvt_vec_t token_scale_vec(a_scale[j]); |
| 259 | cvt_vec_t token_zp_scale_vec; |
| 260 | if constexpr (AZP) { |
| 261 | float zp_scale_val = a_scale[j] * static_cast<float>(azp[j]); |
| 262 | token_zp_scale_vec = cvt_vec_t(zp_scale_val); |
| 263 | } |
| 264 | int64_t k = start; |
| 265 | const float* input_ptr = input + j * hidden_size; |
| 266 | scalar_t* output_ptr = output + j * hidden_size; |
| 267 | for (; k < end - vec_elem_num; k += vec_elem_num) { |
| 268 | cvt_vec_t elems_fp32(input_ptr + k); |
| 269 | elems_fp32 = elems_fp32 * token_scale_vec; |
| 270 | if constexpr (AZP) { |
| 271 | cvt_vec_t azp_adj_fp32(azp_adj + k); |
| 272 | elems_fp32 = elems_fp32 - azp_adj_fp32 * token_zp_scale_vec; |
| 273 | } |
| 274 | if constexpr (Bias) { |
| 275 | load_vec_t bias_vec(bias + k); |
| 276 | cvt_vec_t bias_vec_fp32(bias_vec); |
| 277 | elems_fp32 = elems_fp32 + bias_vec_fp32; |
| 278 | } |
| 279 | load_vec_t elems_out(elems_fp32); |
| 280 | elems_out.save(output_ptr + k); |
| 281 | } |
| 282 | if (k < end) { |
| 283 | cvt_vec_t elems_fp32(input_ptr + k); |
| 284 | elems_fp32 = elems_fp32 * token_scale_vec; |
| 285 | if constexpr (AZP) { |
| 286 | cvt_vec_t azp_adj_fp32(azp_adj + k); |
| 287 | elems_fp32 = elems_fp32 - azp_adj_fp32 * token_zp_scale_vec; |
| 288 | } |
| 289 | if constexpr (Bias) { |
| 290 | load_vec_t bias_vec(bias + k); |
| 291 | cvt_vec_t bias_vec_fp32(bias_vec); |
| 292 | elems_fp32 = elems_fp32 + bias_vec_fp32; |
| 293 | } |
| 294 | load_vec_t elems_out(elems_fp32); |
| 295 | elems_out.save(output_ptr + k, end - k); |
| 296 | } |
| 297 | } |
| 298 | } |
| 299 | } |
| 300 | } |
| 301 | } // namespace |
| 302 | |
| 303 | int64_t create_onednn_scaled_mm_handler( |
| 304 | const torch::Tensor& b, // [IC, OC], column-major |
| 305 | const torch::Tensor& b_scales, // [1] or [OC] |
| 306 | at::ScalarType output_type, bool dynamic_act_quant, bool use_azp, |
| 307 | int64_t primitive_cache_size) { |
| 308 | TORCH_CHECK(b.dim() == 2); |
| 309 | TORCH_CHECK(b.stride(0) == 1); // Column-major |
| 310 | TORCH_CHECK(b_scales.is_contiguous()); |
| 311 | |
| 312 | W8A8MatMulPrimitiveHandler::Args args; |
| 313 | args.primitive_cache_size = primitive_cache_size; |
| 314 | |
| 315 | if (b_scales.numel() == 1) { |
| 316 | args.b_quantization_strategy = |
| 317 | W8A8MatMulPrimitiveHandler::QuantizationStrategy::PER_TENSOR; |
| 318 | } else { |
| 319 | TORCH_CHECK_EQ(b_scales.numel(), b.size(1)); |
| 320 | args.b_quantization_strategy = |
| 321 | W8A8MatMulPrimitiveHandler::QuantizationStrategy::PER_OUTPUT_CHANNEL; |
| 322 | } |
| 323 | args.b_scales_ptr = b_scales.data_ptr<float>(); |
| 324 | args.b_k_size = b.size(0); |
| 325 | args.b_k_stride = b.stride(0); |
| 326 | args.b_n_size = b.size(1); |
| 327 | args.b_n_stride = b.stride(1); |
| 328 | args.b_ptr = b.data_ptr<int8_t>(); |
| 329 | |
| 330 | if (dynamic_act_quant) { |
| 331 | // dynamic per-token, bias, A scales and A zps will be applied in outside. |
| 332 | args.a_quantization_strategy = |
| 333 | W8A8MatMulPrimitiveHandler::QuantizationStrategy::PER_TOKEN; |
| 334 | args.use_a_zero_point = false; |
| 335 | } else { |
| 336 | // static per-tensor |
| 337 | args.a_quantization_strategy = |
| 338 | W8A8MatMulPrimitiveHandler::QuantizationStrategy::PER_TENSOR; |
| 339 | args.use_a_zero_point = use_azp; |
| 340 | } |
| 341 | |
| 342 | VLLM_DISPATCH_FLOATING_TYPES(output_type, "create_onednn_scaled_mm_handler", |
| 343 | [&] { |
| 344 | if (dynamic_act_quant) { |
| 345 | args.c_type = get_dnnl_type<float>(); |
| 346 | } else { |
| 347 | args.c_type = get_dnnl_type<scalar_t>(); |
| 348 | } |
| 349 | }); |
| 350 | |
| 351 | return reinterpret_cast<int64_t>(new W8A8MatMulPrimitiveHandler(args)); |
| 352 | } |
| 353 | |
| 354 | void onednn_scaled_mm( |
| 355 | torch::Tensor& c, // [M, OC], row-major |
| 356 | const torch::Tensor& a, // [M, IC], row-major |
| 357 | const torch::Tensor& a_scales, // [M] or [1] |
| 358 | const std::optional<torch::Tensor>& azp, // [M] or [1] |
| 359 | const std::optional<torch::Tensor>& azp_adj, // [M] or [1] |
| 360 | const std::optional<torch::Tensor>& bias, // [N] |
| 361 | const torch::Tensor& handler_tensor) { |
| 362 | CPU_KERNEL_GUARD_IN(onednn_scaled_mm) |
| 363 | TORCH_CHECK(a.dim() == 2); |
| 364 | TORCH_CHECK(a.is_contiguous()); |
| 365 | TORCH_CHECK(c.is_contiguous()); |
| 366 | W8A8MatMulPrimitiveHandler* ptr = |
| 367 | reinterpret_cast<W8A8MatMulPrimitiveHandler*>( |
| 368 | handler_tensor.item<int64_t>()); |
| 369 | const int32_t* azp_ptr = nullptr; |
| 370 | if (azp.has_value()) { |
| 371 | azp_ptr = azp->data_ptr<int32_t>(); |
| 372 | } |
| 373 | if (ptr->get_input_scale_strategy() == |
| 374 | W8A8MatMulPrimitiveHandler::QuantizationStrategy::PER_TENSOR) { |
| 375 | TORCH_CHECK_EQ(a_scales.numel(), 1); |
| 376 | } |
| 377 | |
| 378 | W8A8MatMulPrimitiveHandler::ExecArgs exec_args; |
| 379 | exec_args.a_ptr = a.data_ptr<int8_t>(); |
| 380 | exec_args.a_m_size = a.size(0); |
| 381 | exec_args.bias_ptr = nullptr; |
| 382 | exec_args.bias_type = get_dnnl_type<void>(); |
| 383 | exec_args.use_bias = false; |
| 384 | exec_args.a_scales_ptr = nullptr; |
| 385 | exec_args.a_zero_points_ptr = nullptr; |
| 386 | |
| 387 | VLLM_DISPATCH_FLOATING_TYPES(c.scalar_type(), "onednn_scaled_mm", [&] { |
| 388 | if (ptr->get_input_scale_strategy() == |
| 389 | W8A8MatMulPrimitiveHandler::QuantizationStrategy::PER_TENSOR) { |
| 390 | if (bias.has_value()) { |
| 391 | exec_args.bias_ptr = bias->data_ptr<scalar_t>(); |
| 392 | exec_args.bias_type = get_dnnl_type<scalar_t>(); |
| 393 | exec_args.use_bias = true; |
| 394 | } |
| 395 | exec_args.a_scales_ptr = a_scales.data_ptr<float>(); |
| 396 | exec_args.a_zero_points_ptr = azp_ptr; |
| 397 | exec_args.c_ptr = c.data_ptr<scalar_t>(); |
| 398 | ptr->execute(exec_args); |
| 399 | } else if (ptr->get_input_scale_strategy() == |
| 400 | W8A8MatMulPrimitiveHandler::QuantizationStrategy::PER_TOKEN) { |
| 401 | torch::Tensor tmp_fp32_out = |
| 402 | torch::empty_like(c, ::at::ScalarType::Float); |
| 403 | exec_args.c_ptr = tmp_fp32_out.data_ptr<float>(); |
| 404 | ptr->execute(exec_args); |
| 405 | if (bias.has_value()) { |
| 406 | if (azp.has_value()) { |
| 407 | dynamic_quant_epilogue<true, true>( |
| 408 | tmp_fp32_out.data_ptr<float>(), c.data_ptr<scalar_t>(), |
| 409 | a_scales.data_ptr<float>(), azp_ptr, azp_adj->data_ptr<float>(), |
| 410 | bias->data_ptr<scalar_t>(), c.size(0), c.size(1)); |
| 411 | } else { |
| 412 | dynamic_quant_epilogue<false, true>( |
| 413 | tmp_fp32_out.data_ptr<float>(), c.data_ptr<scalar_t>(), |
| 414 | a_scales.data_ptr<float>(), azp_ptr, nullptr, |
| 415 | bias->data_ptr<scalar_t>(), c.size(0), c.size(1)); |
| 416 | } |
| 417 | } else { |
| 418 | if (azp.has_value()) { |
| 419 | dynamic_quant_epilogue<true, false>( |
| 420 | tmp_fp32_out.data_ptr<float>(), c.data_ptr<scalar_t>(), |
| 421 | a_scales.data_ptr<float>(), azp_ptr, azp_adj->data_ptr<float>(), |
| 422 | (scalar_t*)nullptr, c.size(0), c.size(1)); |
| 423 | } else { |
| 424 | dynamic_quant_epilogue<false, false>( |
| 425 | tmp_fp32_out.data_ptr<float>(), c.data_ptr<scalar_t>(), |
| 426 | a_scales.data_ptr<float>(), azp_ptr, nullptr, (scalar_t*)nullptr, |
| 427 | c.size(0), c.size(1)); |
| 428 | } |
| 429 | } |
| 430 | } else { |
| 431 | TORCH_CHECK(false, "invalid act quant type."); |
| 432 | } |
| 433 | }); |
| 434 | } |
| 435 | |
| 436 | // static-per-tensor quantization. |
| 437 | void static_scaled_int8_quant( |
| 438 | torch::Tensor& out, // [batch, hidden_size] |
| 439 | const torch::Tensor& input, // [batch, hidden_size] |
| 440 | const torch::Tensor& scale, std::optional<torch::Tensor> const& azp) { |
| 441 | CPU_KERNEL_GUARD_IN(static_scaled_int8_quant) |
| 442 | TORCH_CHECK(out.is_contiguous()); |
| 443 | TORCH_CHECK_EQ(input.dim(), 2); |
| 444 | TORCH_CHECK_EQ(input.stride(1), 1); |
| 445 | TORCH_CHECK(scale.numel() == 1); |
| 446 | TORCH_CHECK(!azp.has_value() || azp->numel() == 1); |
| 447 | |
| 448 | const int64_t stride = input.stride(0); |
| 449 | const int64_t hidden_size = input.size(1); |
| 450 | const int64_t num_tokens = input.size(0); |
| 451 | VLLM_DISPATCH_FLOATING_TYPES( |
| 452 | input.scalar_type(), "static_scaled_int8_quant_impl", [&] { |
| 453 | if (azp.has_value()) { |
| 454 | static_scaled_int8_quant_impl<true>( |
| 455 | input.data_ptr<scalar_t>(), out.data_ptr<int8_t>(), |
| 456 | scale.data_ptr<float>(), azp->data_ptr<int32_t>(), num_tokens, |
| 457 | stride, hidden_size); |
| 458 | } else { |
| 459 | static_scaled_int8_quant_impl<false>(input.data_ptr<scalar_t>(), |
| 460 | out.data_ptr<int8_t>(), |
| 461 | scale.data_ptr<float>(), nullptr, |
| 462 | num_tokens, stride, hidden_size); |
| 463 | } |
| 464 | }); |
| 465 | } |
| 466 | |
| 467 | // dynamic-per-token quantization. |
| 468 | void dynamic_scaled_int8_quant( |
| 469 | torch::Tensor& out, // [batch, hidden_size] |
| 470 | const torch::Tensor& input, // [batch, hidden_size] |
| 471 | torch::Tensor& scale, // [batch, 1] |
| 472 | std::optional<torch::Tensor> const& azp) { |
| 473 | CPU_KERNEL_GUARD_IN(dynamic_scaled_int8_quant) |
| 474 | TORCH_CHECK(out.is_contiguous()); |
| 475 | TORCH_CHECK_EQ(input.dim(), 2); |
| 476 | TORCH_CHECK_EQ(input.stride(1), 1); |
| 477 | |
| 478 | const int64_t hidden_size = input.size(1); |
| 479 | const int64_t num_tokens = input.size(0); |
| 480 | const int64_t stride = input.stride(0); |
| 481 | VLLM_DISPATCH_FLOATING_TYPES( |
| 482 | input.scalar_type(), "dynamic_scaled_int8_quant_impl", [&] { |
| 483 | if (azp.has_value()) { |
| 484 | dynamic_scaled_int8_quant_impl<true>( |
| 485 | input.data_ptr<scalar_t>(), out.data_ptr<int8_t>(), |
| 486 | scale.data_ptr<float>(), azp->data_ptr<int32_t>(), num_tokens, |
| 487 | stride, hidden_size); |
| 488 | } else { |
| 489 | dynamic_scaled_int8_quant_impl<false>( |
| 490 | input.data_ptr<scalar_t>(), out.data_ptr<int8_t>(), |
| 491 | scale.data_ptr<float>(), nullptr, num_tokens, stride, |
| 492 | hidden_size); |
| 493 | } |
| 494 | }); |
| 495 | } |
| 496 | |
| 497 | int64_t create_onednn_mm_handler(const torch::Tensor& b, |
| 498 | int64_t primitive_cache_size) { |
| 499 | TORCH_CHECK(b.dim() == 2); |
| 500 | |
| 501 | MatMulPrimitiveHandler::Args args; |
| 502 | args.primitive_cache_size = primitive_cache_size; |
| 503 | |
| 504 | args.b_k_size = b.size(0); |
| 505 | args.b_k_stride = b.stride(0); |
| 506 | args.b_n_size = b.size(1); |
| 507 | args.b_n_stride = b.stride(1); |
| 508 | args.b_ptr = b.data_ptr(); |
| 509 | |
| 510 | VLLM_DISPATCH_FLOATING_TYPES(b.scalar_type(), "create_onednn_mm_handler", |
| 511 | [&] { |
| 512 | args.c_type = get_dnnl_type<scalar_t>(); |
| 513 | args.ab_type = get_dnnl_type<scalar_t>(); |
| 514 | }); |
| 515 | |
| 516 | #ifdef VLLM_USE_ACL |
| 517 | // TODO(fadara01): this is a hot-fix, remove once problem is addressed in the |
| 518 | // AArch64 stack. hybrid kernels in ACL which are the fastest for LM head |
| 519 | // layers are only enabled for BF16 x BF16 -> FP32 problem desc as they can't |
| 520 | // convert FP32->BF16 internally. In this case do: oneDNN matmul with |
| 521 | // bf16:bf16:fp32 then cast output down to bf16. |
| 522 | if (b.scalar_type() == at::kBFloat16 && b.size(1) >= 64000) { |
| 523 | args.c_type = dnnl::memory::data_type::f32; |
| 524 | } |
| 525 | #endif |
| 526 | |
| 527 | return reinterpret_cast<int64_t>(new MatMulPrimitiveHandler(args)); |
| 528 | } |
| 529 | |
| 530 | void onednn_mm(torch::Tensor& c, // [M, OC], row-major |
| 531 | const torch::Tensor& a, // [M, IC], row-major |
| 532 | const std::optional<torch::Tensor>& bias, |
| 533 | const torch::Tensor& handler_tensor) { |
| 534 | CPU_KERNEL_GUARD_IN(onednn_mm) |
| 535 | TORCH_CHECK(a.dim() == 2); |
| 536 | TORCH_CHECK(a.stride(-1) == 1); |
| 537 | TORCH_CHECK(c.stride(-1) == 1); |
| 538 | MatMulPrimitiveHandler* ptr = |
| 539 | reinterpret_cast<MatMulPrimitiveHandler*>(handler_tensor.item<int64_t>()); |
| 540 | |
| 541 | torch::Tensor dnnl_output = c; |
| 542 | // AArch64 specific case where we do a bf16 x bf16 -> fp32 matmul |
| 543 | const bool convert_output = |
| 544 | ptr->get_output_type() == dnnl::memory::data_type::f32 && |
| 545 | c.scalar_type() != at::kFloat; |
| 546 | if (convert_output) { |
| 547 | dnnl_output = torch::empty(c.sizes(), c.options().dtype(at::kFloat)); |
| 548 | } |
| 549 | |
| 550 | // ACL matmuls expect contiguous source tensors |
| 551 | #ifdef VLLM_USE_ACL |
| 552 | torch::Tensor a_contig = a.contiguous(); |
| 553 | #endif |
| 554 | |
| 555 | MatMulPrimitiveHandler::ExecArgs exec_args; |
| 556 | exec_args.c_ptr = dnnl_output.data_ptr(); |
| 557 | |
| 558 | #ifdef VLLM_USE_ACL |
| 559 | exec_args.a_m_size = a_contig.size(0); |
| 560 | exec_args.a_m_stride = a_contig.stride(0); |
| 561 | #else |
| 562 | exec_args.a_m_size = a.size(0); |
| 563 | exec_args.a_m_stride = a.stride(0); |
| 564 | #endif |
| 565 | VLLM_DISPATCH_FLOATING_TYPES(a.scalar_type(), "onednn_mm", [&] { |
| 566 | if (bias.has_value()) { |
| 567 | exec_args.use_bias = true; |
| 568 | exec_args.bias_type = get_dnnl_type<scalar_t>(); |
| 569 | #ifdef VLLM_USE_ACL |
| 570 | // ACL matmuls in oneDNN do not support a bias. |
| 571 | // We handle a matmul with bias by doing: c = bias; c += matmul(a, b) |
| 572 | dnnl_output.copy_(bias.value()); |
| 573 | #else |
| 574 | exec_args.bias_ptr = bias->data_ptr<scalar_t>(); |
| 575 | #endif |
| 576 | } else { |
| 577 | exec_args.use_bias = false; |
| 578 | exec_args.bias_type = get_dnnl_type<void>(); |
| 579 | exec_args.bias_ptr = nullptr; |
| 580 | } |
| 581 | #ifdef VLLM_USE_ACL |
| 582 | exec_args.a_ptr = a_contig.data_ptr<scalar_t>(); |
| 583 | #else |
| 584 | exec_args.a_ptr = a.data_ptr<scalar_t>(); |
| 585 | |
| 586 | #endif |
| 587 | |
| 588 | ptr->execute(exec_args); |
| 589 | }); |
| 590 | |
| 591 | if (convert_output) { |
| 592 | c.copy_(dnnl_output); |
| 593 | } |
| 594 | } |