[inductor][cpp] Add FlexAttention support for CPU inference (pytorch#141453)

This PR brings the FlexAttention inference support for the inductor backend in torch.compile (support precisions: bf16 and fp32) on CPUs.

Based on the existing CPP template, this PR extends and implements a FlexAttention CPP template to support broad attention variants, and meanwhile brings optimized performance on CPUs.

With this, users can transparently extend their Flex Attention usages to CPUs with good and common support from torch.compile, both functionality and performance.

For UT tests, in this PR, we include partial critical tests for CPUs as the following (conduct inference tests):
```
pytest test/inductor/test_flex_attention.py
`TestFlexAttention`
#common functions:
run_test
preprocess_paged_attention
run_paged_attention
run_test_with_paged_attention
run_test_with_call
run_dynamic_test
run_automatic_dynamic_test

#test functions:
test_builtin_score_mods
test_builtin_score_mods_automatic_dynamic
test_builtin_score_mods_different_seqlen
test_builtin_score_mods_different_block_size
test_kv_batch_broadcast
test_GQA
test_cpu_error_message_return_lse
test_validate_cpu_dtype_error_message

`TestPagedAttention`
#test function:
test_paged_builtin_score_mods
```
For the rest UTs in `test/inductor/test_flex_attention.py ` and `test/inductor/test_flex_decoding.py`, due to bigger lines of changes (1500+ LOC) that make this PR hard to review, will submit another PR specific for CPU device UTs enabling and refactor.

Besides, more optimizations are also planned in follow up PRs, including:

- Block sparse computation
- Flash decoding tuning

Pull Request resolved: pytorch#141453
Approved by: https://github.com/drisspg, https://github.com/leslie-fang-intel

Co-authored-by: Wu, Chunyuan <[email protected]>

Author: 2 people

aten/src/ATen/native/CPUBlas.cpp

@@ -1125,6 +1125,9 @@ struct Brgemm : public KernelCache <BrgemmKey, GemmHelper> {
     if (dtype == ScalarType::Half) {
       static bool fp16_support = dnnl::get_effective_cpu_isa() >= dnnl::cpu_isa::avx512_core_fp16;
       return fp16_support;
+    } else if (dtype == ScalarType::Float) {
+      static bool fp32_support = dnnl::get_effective_cpu_isa() >= dnnl::cpu_isa::avx2;
+      return fp32_support;
     } else if (dtype == ScalarType::BFloat16) {
       static bool bf16_support = dnnl::get_effective_cpu_isa() >= dnnl::cpu_isa::avx512_core;
       return bf16_support;
@@ -1192,18 +1195,29 @@ void brgemm(
     int64_t ld_b,
     int64_t ld_c,
     const bool add_C,
-    const at::Half* A,
-    const at::Half* B,
-    float* C) {
+    const float* A,
+    const float* B,
+    float* C,
+    bool is_vnni) {
+
+  TORCH_CHECK(!is_vnni,
+    "Float Brgemm does not support vnni layout.");
+
 #if defined(ONEDNN_UKERNEL_ENABLED)
-  if (Brgemm::device_check(ScalarType::Half)) {
-    Brgemm::call<at::Half, at::Half, float>(
+  if (Brgemm::device_check(ScalarType::Float)) {
+    Brgemm::call<float, float, float>(
       M, N, K, ld_a, ld_b, ld_c, add_C, A, B, C);
     return;
   }
 #endif
-  TORCH_CHECK(false,
-  "Half Brgemm is only supported on X64 when oneDNN ukernel is enabled and avx512_fp16 is supported");
+  // fallback path
+  auto beta = add_C ? 1 : 0;
+  gemm(
+    at::native::TransposeType::NoTranspose,
+    at::native::TransposeType::NoTranspose,
+    N, M, K, 1,
+    B, ld_b, A, ld_a,
+    beta, C, ld_c);
 }
 
 void brgemm(
@@ -1216,22 +1230,64 @@ void brgemm(
     const bool add_C,
     const at::BFloat16* A,
     const at::BFloat16* B,
-    float* C) {
+    float* C,
+    bool is_vnni) {
 #if defined(ONEDNN_UKERNEL_ENABLED)
-  if (Brgemm::device_check(ScalarType::BFloat16)) {
+  if (is_vnni && Brgemm::device_check(ScalarType::BFloat16)) {
     Brgemm::call<at::BFloat16, at::BFloat16, float>(
       M, N, K, ld_a, ld_b, ld_c, add_C, A, B, C);
     return;
   }
 #endif
-  TORCH_CHECK(false,
-  "BFloat16 Brgemm is only supported on X64 when oneDNN ukernel is enabled and avx512 is supported");
+  // fallback path
+  TORCH_CHECK(!is_vnni,
+    "BFloat16 Brgemm VNNI format is only supported on X64 when oneDNN ukernel is enabled and `amx` is supported");
+  auto beta = add_C ? 1 : 0;
+  gemm(
+    at::native::TransposeType::NoTranspose,
+    at::native::TransposeType::NoTranspose,
+    N, M, K, 1,
+    B, ld_b, A, ld_a,
+    beta, C, ld_c);
+}
+
+void brgemm(
+    int64_t M,
+    int64_t N,
+    int64_t K,
+    int64_t ld_a,
+    int64_t ld_b,
+    int64_t ld_c,
+    const bool add_C,
+    const at::Half* A,
+    const at::Half* B,
+    float* C,
+    bool is_vnni) {
+#if defined(ONEDNN_UKERNEL_ENABLED)
+  if (is_vnni && Brgemm::device_check(ScalarType::Half)) {
+    Brgemm::call<at::Half, at::Half, float>(
+      M, N, K, ld_a, ld_b, ld_c, add_C, A, B, C);
+    return;
+  }
+#endif
+  // fallback path
+  TORCH_CHECK(!is_vnni,
+    "Half Brgemm VNNI format is only supported on X64 when oneDNN ukernel is enabled and `amx_fp16` is supported");
+  auto beta = add_C ? 1 : 0;
+  gemm(
+    at::native::TransposeType::NoTranspose,
+    at::native::TransposeType::NoTranspose,
+    N, M, K, 1,
+    B, ld_b, A, ld_a,
+    beta, C, ld_c);
 }
 
-void brgemm_release() {
+void brgemm_release(bool is_vnni) {
 #if defined(ONEDNN_UKERNEL_ENABLED)
-  dnnl::ukernel::brgemm::release_hw_context();
-  Brgemm::get_current() = nullptr;
+  if (is_vnni) {
+    dnnl::ukernel::brgemm::release_hw_context();
+    Brgemm::get_current() = nullptr;
+  }
 #endif
 }
 

aten/src/ATen/native/CPUBlas.h

@@ -204,7 +204,8 @@ TORCH_API void brgemm(
     const bool add_C,
     const at::Half* A,
     const at::Half* B,
-    float* C);
+    float* C,
+    bool is_vnni = true);
 
 TORCH_API void brgemm(
     int64_t M,
@@ -216,10 +217,24 @@ TORCH_API void brgemm(
     const bool add_C,
     const at::BFloat16* A,
     const at::BFloat16* B,
-    float* C);
+    float* C,
+    bool is_vnni = true);
+
+TORCH_API void brgemm(
+    int64_t M,
+    int64_t N,
+    int64_t K,
+    int64_t ld_a,
+    int64_t ld_b,
+    int64_t ld_c,
+    const bool add_C,
+    const float* A,
+    const float* B,
+    float* C,
+    bool is_vnni = false);
 
 // Release brgemm hardware context
-TORCH_API void brgemm_release();
+TORCH_API void brgemm_release(bool is_vnni = true);
 
 // Pack B matrix to get better performance if needed
 void pack(
@@ -233,6 +248,6 @@ void pack(
     void* out);
 
 // Whether pack is supported in the platform.
-bool could_pack(ScalarType dt_in);
+TORCH_API bool could_pack(ScalarType dt_in);
 
 } // namespace at::native::cpublas