执行摘要

• 一句话：为AVX2 CPU添加DNNL W8A8 INT8量化支持
• 推荐动作：值得精读。重点关注 CMake 中对多 ISA 库的链接策略（统一 dnnl_ext 而非分离编译）以及 C++ 向量类型中条件编译 fallback 的模式。该设计在保持代码复用性的同时提供了向前兼容。建议后续补充自动化测试，覆盖 AVX2 平台的量化算子。

功能与动机

CPU 后端的 W8A8 INT8 量化操作（static_scaled_int8_quant、dynamic_scaled_int8_quant、onednn_scaled_mm）被 AVX512F 守卫，在仅支持 AVX2 的主机（如 Xeon-6 E-core）上运行压缩张量模型会触发运行时符号缺失错误。此外，int8 量化对 AVX2 尤其有益，因为 bf16/fp16 模型在 AVX2 上只能以 fp32 速率运行。

实现拆解

1. 构建系统调整（cmake/cpu_extension.cmake）：移除旧的 DNNL 编译标志选择逻辑，改为在 x86 平台统一使用 AVX2 标志编译 dnnl_ext（-mavx2），使得同一份 dnnl_ext 库同时服务于 AVX2 和 AVX512。将 _C_AVX2 的链接库从仅 numa 扩展到 numa dnnl_ext，并在源码列表中加入 dnnl_kernels.cpp 和 torch_bindings.cpp。
2. 算子注册扩展（csrc/cpu/torch_bindings.cpp）：将 #if 条件中的 __AVX512F__ 扩展为 __AVX512F__ || __AVX2__，使 onednn_mm、onednn_scaled_mm 等量化算子对 AVX2 可见。
3. 向量类型补齐（csrc/cpu/cpu_types_x86.hpp）：为 FP16Vec16、BF16Vec16 的 save(void* ptr, int elem_num) 添加 #else 分支——当缺少 __AVX512BW__ 时使用临时数组配合标量循环实现 partial store。为 FP32Vec16 新增取反运算符（operator-）和 partial store 方法（使用 _mm256_maskstore_ps 或条件 store）。为 BF16Vec32 添加默认构造函数和简化版 BF16Vec8 扩展构造函数（改用 _mm256_broadcastsi128_si256）。
4. 量化内核 bugfix（csrc/cpu/dnnl_kernels.cpp）：在 dynamic_quant_epilogue 模板函数中，将主循环增量从 ++j 改为 j += vec_elem_num，避免重复处理同一元素导致错误。
5. 测试与验证：未新增自动化测试文件，但 PR body 提供了在 AVX2 平台上的端到端压测结果（50 个 prompt，128 输入/输出长度），并附有 int8 与 bf16 的吞吐量对比。

关键文件：

• csrc/cpu/cpu_types_x86.hpp（模块 CPU向量类型；类别 source；类型 core-logic；符号 FP16Vec16::save, BF16Vec16::save, FP32Vec16::save, FP32Vec16::operator-）: 核心向量类型扩展，为 AVX2 添加了 FP16Vec16/BF16Vec16 partial store fallback、FP32Vec16 取反和 partial store、BF16Vec32 默认构造器和简化版 BF16Vec8 构造器等，是本次变更中改动量最大的文件。
• cmake/cpu_extension.cmake（模块 构建配置；类别 other；类型 core-logic）: 构建系统的核心变更：移除旧的 DNNL 编译标志选择，统一以 AVX2 标志编译 dnnl_ext；将 _C_AVX2 链接到 dnnl_ext 并添加 dnnl_kernels.cpp 和 torch_bindings.cpp。
• csrc/cpu/torch_bindings.cpp（模块 算子绑定；类别 source；类型 core-logic）: 修改量化算子的编译条件 gate，新增 AVX2 使 onednn 相关算子对 AVX2 平台可见。
• csrc/cpu/dnnl_kernels.cpp（模块 量化内核；类别 source；类型 core-logic；符号 dynamic_quant_epilogue）: 修复 dynamic_quant_epilogue 模板函数中的循环步长 bug，该 bug 导致在 AVX2 路径下同一元素会被重复处理。

关键符号：FP16Vec16::save, BF16Vec16::save, FP32Vec16::save, FP32Vec16::operator-, BF16Vec32::BF16Vec32, dynamic_quant_epilogue

关键源码片段

csrc/cpu/cpu_types_x86.hpp

核心向量类型扩展，为 AVX2 添加了 FP16Vec16/BF16Vec16 partial store fallback、FP32Vec16 取反和 partial store、BF16Vec32 默认构造器和简化版 BF16Vec8 构造器等，是本次变更中改动量最大的文件。

// FP16Vec16 partial store: 使用 AVX512BW 的 mask store 或在缺少时回退到临时数组
void save(void* ptr, const int elem_num) const {
#ifdef __AVX512BW__
    constexpr uint32_t M = 0xFFFFFFFF;
    __mmask16 mask = _cvtu32_mask16(M >> (32 - elem_num));
    _mm256_mask_storeu_epi16(ptr, mask, reg);
#else
    // 软件 fallback：先 store 整个向量，再逐个写入有效元素
    int16_t tmp[VEC_ELEM_NUM];
    _mm256_storeu_si256((__m256i*)tmp, reg);
    for (int i = 0; i < elem_num; ++i)
      reinterpret_cast<int16_t*>(ptr)[i] = tmp[i];
#endif
}
​
// BF16Vec32 默认构造：初始化为零
BF16Vec32()
    : reg_low(_mm256_setzero_si256()), reg_high(_mm256_setzero_si256()) {}
​
// BF16Vec32 从 BF16Vec8 构造：使用 broadcast 代替多次 insert
BF16Vec32(BF16Vec8& vec8_data)
    : reg_low(_mm256_broadcastsi128_si256((__m128i)vec8_data.reg)),
      reg_high(_mm256_broadcastsi128_si256((__m128i)vec8_data.reg)) {}

cmake/cpu_extension.cmake

构建系统的核心变更：移除旧的 DNNL 编译标志选择，统一以 AVX2 标志编译 dnnl_ext；将 _C_AVX2 链接到 dnnl_ext 并添加 dnnl_kernels.cpp 和 torch_bindings.cpp。

# 移除旧的 AVX512 专用 DNNL 编译标志，统一使用 AVX2
# before:
# if (ENABLE_X86_ISA)
# list(APPEND DNNL_COMPILE_FLAGS ${CXX_COMPILE_FLAGS_AVX512})
# else()
# list(APPEND DNNL_COMPILE_FLAGS ${CXX_COMPILE_FLAGS})
# endif()
# after:
if (ENABLE_X86_ISA)
  target_compile_options(dnnl_ext PRIVATE ${CXX_COMPILE_FLAGS_AVX2} -fPIC)
else()
  target_compile_options(dnnl_ext PRIVATE ${CXX_COMPILE_FLAGS} -fPIC)
endif()
​
# 将 _C_AVX2 与 dnnl_ext 链接
set(_C_AVX2_LIBS numa dnnl_ext)

csrc/cpu/dnnl_kernels.cpp

修复 dynamic_quant_epilogue 模板函数中的循环步长 bug，该 bug 导致在 AVX2 路径下同一元素会被重复处理。

// 修复前：j 每次递增 1，导致对同一个元素重复量化
// for (; j < hidden_size - vec_elem_num; ++j) {
// 修复后：每次跳过 vec_elem_num 个元素
for (; j < hidden_size - vec_elem_num; j += vec_elem_num) {
    cvt_vec_t elems_fp32(input_ptr + j);
    // ... 量化逻辑
}

评论区精华

作者在 issue 评论中提到 Apple Silicon 的 smoke test 失败（需合并 #41387 后 rebase），同时指出 ARM 平台使用 DNNL 时可能存在编译问题。审核人 bigPYJ1151 直接批准（LGTM），无深入的设计争辩。自动 bot（Claude Code Review 和 Gemini Code Assist）未提出反对意见。

• Apple Silicon 和 ARM DNNL 兼容性 (other): 未在 PR 内解决，作者计划在后续 PR 中修复 ARM 编译问题。当前 PR 专注于 AVX2，非 x86 平台行为保持不变。

风险与影响

• 风险：

  1. 性能风险（csrc/cpu/cpu_types_x86.hpp）：AVX2 上 partial store 使用标量回退（临时数组 + for 循环），相比 AVX512 的 masking 指令可能带来少量开销，但仅发生在向量不满的场景（通常为尾块），影响有限。
  2. ARM/PowerPC 兼容性（cmake/cpu_extension.cmake、csrc/cpu/torch_bindings.cpp）：编译条件中对非 x86 平台保留了原有行为，但 dnnl_ext 的编译标志在非 x86 时使用通用 CXX_COMPILE_FLAGS，可能引入未预期的编译错误（author 已提及 ARM 问题）。
  3. 缺少单元测试：本次变更未添加专用测试文件，仅依赖手动压测。若未来重构向量类型或 CMake 逻辑，容易引入回归。
  4. oneDNN runtime JIT：oneDNN 在运行时根据 ISA 动态 JIT 编译内核，但本次将整个 dnnl_ext 以 -mavx2 编译，确保了一致性，预期安全。
     - 影响：用户角度：AVX2-only CPU（如部分 Xeon 6 E-core、旧款台式机）的用户现在可以正确加载并运行 W8A8 INT8 量化的模型，int8 吞吐量相比 bf16 提升约 2.4 倍。此前此类用户只能退避到 fp32 或无法使用量化模型。
     系统角度：CMake 构建脚本统一了 dnnl_ext 的编译方式（不再为 AVX512 单独设置 flags），未来维护时需注意 x86 与非 x86 的路径分支。
     团队角度：需要维护两套 vector 操作的 ISA 分支（现有 AVX512 和新添 AVX2 fallback），但核心量化内核共享同一份模板。

• 风险标记：AVX2 fallback 性能路径, ARM/PowerPC 编译待验证, 缺少自动化测试覆盖

关联脉络

• PR #41387 [Fix] Add missing stubs from cpu fp8 attention changes: 作者明确提到该 PR 需要在 #41387 合并后 rebase，因为 Apple Silicon 的 smoke test 依赖于 #41387 中的存根修正。