diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index cca2a98cee57..414fbc5e238c 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -173,6 +173,18 @@ if (CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64" OR ENABLE_X86_ISA) set(AMX_FP8_SUPPORTED FALSE) message(STATUS "AMX-FP8 disabled by ENABLE_AMX_FP8=OFF") endif() + + # AVX10.2 arrived in GCC 15, so the target attribute, the fp8 convert + # intrinsics and __builtin_cpu_supports("avx10.2") are all hard errors on + # the 12.3+ the backend otherwise supports. Probe instead of assuming. + check_cxx_compiler_flag("-mavx10.2" COMPILER_SUPPORTS_AVX10_2_FLAG) + if (COMPILER_SUPPORTS_AVX10_2_FLAG) + set(AVX10_2_SUPPORTED TRUE) + else() + set(AVX10_2_SUPPORTED FALSE) + message(STATUS "AVX10.2 disabled: compiler does not support -mavx10.2 (compiler: ${CMAKE_CXX_COMPILER_ID} ${CMAKE_CXX_COMPILER_VERSION})") + endif() + list(APPEND CXX_COMPILE_FLAGS_AVX2 "-mavx2") elseif (POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND) @@ -629,6 +641,12 @@ if (ENABLE_X86_ISA) message(STATUS "AMX-FP8 (Diamond Rapids) enabled") endif() + # For the sgl-kernels AVX10.2 fp8 quantize paths + if (AVX10_2_SUPPORTED) + target_compile_definitions(_C PRIVATE "-DCPU_CAPABILITY_AVX10_2") + message(STATUS "AVX10.2 enabled") + endif() + # AVX512F define_extension_target( _C_AVX512 diff --git a/csrc/cpu/sgl-kernels/common.h b/csrc/cpu/sgl-kernels/common.h index dc0769450e36..acf75912819f 100644 --- a/csrc/cpu/sgl-kernels/common.h +++ b/csrc/cpu/sgl-kernels/common.h @@ -532,10 +532,17 @@ void zero_buffer(T* data, int64_t size) { #endif #if defined(__x86_64__) +// CPU_CAPABILITY_AVX10_2 is defined by cmake only when the compiler accepts +// -mavx10.2; without it the builtin below is not merely absent but a hard +// error, so the AVX512 paths have to stand in. bool avx10_2_available() { +#ifdef CPU_CAPABILITY_AVX10_2 // __builtin_cpu_supports returns the masked feature bit rather than 0/1. static const bool available = __builtin_cpu_supports("avx10.2"); return available; +#else + return false; +#endif } #endif diff --git a/csrc/cpu/sgl-kernels/gemm_fp8_w8a8.cpp b/csrc/cpu/sgl-kernels/gemm_fp8_w8a8.cpp index 854566d682bd..2014a29fa295 100644 --- a/csrc/cpu/sgl-kernels/gemm_fp8_w8a8.cpp +++ b/csrc/cpu/sgl-kernels/gemm_fp8_w8a8.cpp @@ -772,6 +772,7 @@ inline __m128i cvtfp32_fp8e4m3(__m512& src) { return _mm512_cvtepi32_epi8(packed); } +#ifdef CPU_CAPABILITY_AVX10_2 __attribute__((target("avx10.2"))) inline __m128i cvtfp32_fp8e4m3_avx10_2(__m512& src) { __m256i f16_vec = _mm512_cvt_roundps_ph(src, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC); return _mm256_cvtph_hf8(_mm256_castsi256_ph(f16_vec)); @@ -805,6 +806,7 @@ __attribute__((target("avx10.2"))) void quantize_bf16_block_avx10_2( _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + i + 16), fp8_vec2); } } +#endif // CPU_CAPABILITY_AVX10_2 std::tuple _quantize_fp8e4m3_bf16_per_tensor_no_scale(const at::Tensor& t) { constexpr float quant_max = 448.0f; // torch.finfo(torch.float8_e4m3fn).max @@ -861,10 +863,12 @@ std::tuple _quantize_fp8e4m3_bf16_per_tensor_no_scale(co scale_data[c] = scale_val; // Step 3: Scale and clamp using AVX512 (reuse the same loop structure) +#ifdef CPU_CAPABILITY_AVX10_2 if (avx10_2_available()) { quantize_bf16_block_avx10_2(channel_src, quant_dst, elements_per_channel, scale_reciprocal); continue; } +#endif i = 0; const __m512 scale_recip_vec = _mm512_set1_ps(scale_reciprocal); const __m512 quant_max_vec = _mm512_set1_ps(quant_max); @@ -939,10 +943,12 @@ _quantize_fp8e4m3_bf16_per_tensor_with_scale(const at::Tensor& t, at::Tensor& sc const at::BFloat16* src_data = t_bf16.data_ptr() + c * elements_per_channel; at::Float8_e4m3fn* quant_t_data = quant_t.data_ptr() + c * elements_per_channel; +#ifdef CPU_CAPABILITY_AVX10_2 if (avx10_2_available()) { quantize_bf16_block_avx10_2(src_data, quant_t_data, elements_per_channel, scale_reciprocal); continue; } +#endif int64_t i = 0; // Process 32 elements at a time using AVX512