From f5463512c6d432192fb8152819511dbc8261e540 Mon Sep 17 00:00:00 2001 From: Tingqian Li Date: Mon, 10 Jun 2024 13:09:26 +0800 Subject: [PATCH] [CPU] Add MLP & QKV fusion for LLM serving with PageAttention (#24837) ## Details: in throughput oriented use case like vLLM, the computation bound of Fullyconnect layers has a chance to become similar or even bigger than memory-bound when batch-size is big, in such case, some optimization opportunities was exploited in this PR: ### fusing multiple FC to get more computations per memory-read - when split the problem `C=matmul(A, B)` among multi-cores, each core has sub-problem of `subC=matmul(A, subB)`, with shapes: - A `[BM, K]` - subB `[K, BN]` - subC `[BM, BN]` if we ignore subC matrix (since it's being loaded/stored only once during `brgemm` procedure), the ratio of computations per memory-reads is ` BM*K*BN / (BM*K + BN*K) = 1/(1/BM + 1/BN)`, and to get more computations per memory-read, we prefer to increase this ratio, so we'd like BN to be bigger (that's why we merge gate_FC with up_FC, and q_proj with k_proj and v_proj). in `down_proj`, there is no way to increase output channels (BN), we split along K dimension once to reduce the number of splits of N, to keep BN big too. ### prefetch sub-weight matrix to hide memory latency by computation for execution of sub-problem `subC=matmul(A, subB)` on single core, we further decompose it into sub-problems as following, to mix SW-prefetching with AMX computations: ```python for k0 in range(0, K, K_step): k1 = k0 + K_step blkB = subB[k0:k1, :] blkA = A[:, k0:k1] next_blkB = subB[k1:k1+K_step, :] # blkB's shape is [K_step, BN], which is L2 cache-able # inside following AMX based matmul jit kernel # we evenly scatter SW-prefetching instructions for next_blkB subC+=matmul_kernel(blkA, blkB, sw_prefetch = next_blkB) . ``` ### flexible post-ops jit-based post-Ops allows us to: - interleaving fusion of weights of gate & up along N dimension in unit of 16 fp32 elements, so post-op jit-kernel can combine them using activation(like SiLU) while both results are hot in cache. - QKV fusion requires final outputs was stored into 3 independent output memories, special post-op jit kernel can do that easily. ### Tickets: - *ticket-id* --------- Co-authored-by: Luo Cheng --- src/plugins/intel_cpu/CMakeLists.txt | 9 + src/plugins/intel_cpu/src/cpu_types.cpp | 4 + src/plugins/intel_cpu/src/cpu_types.h | 2 + src/plugins/intel_cpu/src/extension.cpp | 4 + .../src/nodes/kernels/x64/mlp_kernel.cpp | 385 ++++++++++++++++++ .../src/nodes/kernels/x64/mlp_kernel.hpp | 264 ++++++++++++ .../src/nodes/kernels/x64/mlp_utils.cpp | 28 ++ .../src/nodes/kernels/x64/mlp_utils.hpp | 19 + src/plugins/intel_cpu/src/nodes/llm_mlp.cpp | 339 +++++++++++++++ src/plugins/intel_cpu/src/nodes/llm_mlp.h | 39 ++ src/plugins/intel_cpu/src/nodes/qkv_proj.cpp | 291 +++++++++++++ src/plugins/intel_cpu/src/nodes/qkv_proj.h | 41 ++ src/plugins/intel_cpu/src/nodes_factory.cpp | 4 + .../cpu_opset/x64/op/llm_mlp.cpp | 59 +++ .../cpu_opset/x64/op/llm_mlp.hpp | 62 +++ .../cpu_opset/x64/op/qkv_proj.cpp | 39 ++ .../cpu_opset/x64/op/qkv_proj.hpp | 34 ++ .../cpu_opset/x64/pass/mlp_fusion.cpp | 132 ++++++ .../cpu_opset/x64/pass/mlp_fusion.hpp | 19 + .../cpu_opset/x64/pass/qkv_proj_fusion.cpp | 123 ++++++ .../cpu_opset/x64/pass/qkv_proj_fusion.hpp | 19 + .../transformation_pipeline.cpp | 36 ++ .../intel_cpu/src/utils/plain_tensor.hpp | 4 + 23 files changed, 1956 insertions(+) create mode 100644 src/plugins/intel_cpu/src/nodes/kernels/x64/mlp_kernel.cpp create mode 100644 src/plugins/intel_cpu/src/nodes/kernels/x64/mlp_kernel.hpp create mode 100644 src/plugins/intel_cpu/src/nodes/kernels/x64/mlp_utils.cpp create mode 100644 src/plugins/intel_cpu/src/nodes/kernels/x64/mlp_utils.hpp create mode 100644 src/plugins/intel_cpu/src/nodes/llm_mlp.cpp create mode 100644 src/plugins/intel_cpu/src/nodes/llm_mlp.h create mode 100644 src/plugins/intel_cpu/src/nodes/qkv_proj.cpp create mode 100644 src/plugins/intel_cpu/src/nodes/qkv_proj.h create mode 100644 src/plugins/intel_cpu/src/transformations/cpu_opset/x64/op/llm_mlp.cpp create mode 100644 src/plugins/intel_cpu/src/transformations/cpu_opset/x64/op/llm_mlp.hpp create mode 100644 src/plugins/intel_cpu/src/transformations/cpu_opset/x64/op/qkv_proj.cpp create mode 100644 src/plugins/intel_cpu/src/transformations/cpu_opset/x64/op/qkv_proj.hpp create mode 100644 src/plugins/intel_cpu/src/transformations/cpu_opset/x64/pass/mlp_fusion.cpp create mode 100644 src/plugins/intel_cpu/src/transformations/cpu_opset/x64/pass/mlp_fusion.hpp create mode 100644 src/plugins/intel_cpu/src/transformations/cpu_opset/x64/pass/qkv_proj_fusion.cpp create mode 100644 src/plugins/intel_cpu/src/transformations/cpu_opset/x64/pass/qkv_proj_fusion.hpp diff --git a/src/plugins/intel_cpu/CMakeLists.txt b/src/plugins/intel_cpu/CMakeLists.txt index a62debac98e..14cd2d2f918 100644 --- a/src/plugins/intel_cpu/CMakeLists.txt +++ b/src/plugins/intel_cpu/CMakeLists.txt @@ -245,6 +245,15 @@ cross_compiled_file(${TARGET_NAME} NAME attn_quantkv paged_attn_quantkv attn_quant_u8 attn_dequant_u8 NAMESPACE ov::Extensions::Cpu::XARCH ) + +cross_compiled_file(${TARGET_NAME} + ARCH AVX512F ANY + src/nodes/kernels/x64/mlp_utils.cpp + API src/nodes/kernels/x64/mlp_utils.hpp + NAME llm_mlp_transpose_epi32_16x16 + NAMESPACE ov::Extensions::Cpu::XARCH +) + # system dependencies must go last target_link_libraries(${TARGET_NAME} PRIVATE openvino::pugixml) ov_set_threading_interface_for(${TARGET_NAME}) diff --git a/src/plugins/intel_cpu/src/cpu_types.cpp b/src/plugins/intel_cpu/src/cpu_types.cpp index 3a6335d0648..130fca6d6eb 100644 --- a/src/plugins/intel_cpu/src/cpu_types.cpp +++ b/src/plugins/intel_cpu/src/cpu_types.cpp @@ -244,6 +244,8 @@ static const TypeToNameMap& get_type_to_name_tbl() { {"CausalMaskPreprocess", Type::CausalMaskPreprocess}, {"EmbeddingBagPacked", Type::EmbeddingBagPacked}, {"EmbeddingBagOffsets", Type::EmbeddingBagOffsets}, + {"LLMMLP", Type::LLMMLP}, + {"QKVProjection", Type::QKVProjection}, }; return type_to_name_tbl; } @@ -367,6 +369,8 @@ std::string NameFromType(const Type type) { CASE(PagedAttention); CASE(RoPE); CASE(CausalMaskPreprocess); + CASE(LLMMLP); + CASE(QKVProjection); CASE(Unknown); } #undef CASE diff --git a/src/plugins/intel_cpu/src/cpu_types.h b/src/plugins/intel_cpu/src/cpu_types.h index 24704d7f7b9..bfb411446c2 100644 --- a/src/plugins/intel_cpu/src/cpu_types.h +++ b/src/plugins/intel_cpu/src/cpu_types.h @@ -124,6 +124,8 @@ enum class Type { PagedAttention, RoPE, CausalMaskPreprocess, + LLMMLP, + QKVProjection, }; enum class Algorithm { diff --git a/src/plugins/intel_cpu/src/extension.cpp b/src/plugins/intel_cpu/src/extension.cpp index 9e496ba5cd8..dc0ed0b9839 100644 --- a/src/plugins/intel_cpu/src/extension.cpp +++ b/src/plugins/intel_cpu/src/extension.cpp @@ -23,6 +23,8 @@ #include "transformations/cpu_opset/common/op/swish_cpu.hpp" #include "transformations/cpu_opset/x64/op/interaction.hpp" #include "transformations/cpu_opset/x64/op/mha.hpp" +#include "transformations/cpu_opset/x64/op/llm_mlp.hpp" +#include "transformations/cpu_opset/x64/op/qkv_proj.hpp" #include "transformations/snippets/x64/op/brgemm_copy_b.hpp" #include "transformations/snippets/x64/op/brgemm_cpu.hpp" #include "transformations/snippets/x64/op/load_convert.hpp" @@ -82,6 +84,8 @@ private: OP_EXTENSION(ov::op::internal::RoPE) \ OP_EXTENSION_X64(ov::intel_cpu::MHANode) \ OP_EXTENSION_X64(ov::intel_cpu::InteractionNode) \ + OP_EXTENSION_X64(ov::intel_cpu::LLMMLPNode) \ + OP_EXTENSION_X64(ov::intel_cpu::QKVProjectionNode) \ OP_EXTENSION_X64(ov::intel_cpu::ScaledDotProductAttentionWithKVCache) \ OP_EXTENSION_X64(ov::intel_cpu::LoadConvertSaturation) \ OP_EXTENSION_X64(ov::intel_cpu::LoadConvertTruncation) \ diff --git a/src/plugins/intel_cpu/src/nodes/kernels/x64/mlp_kernel.cpp b/src/plugins/intel_cpu/src/nodes/kernels/x64/mlp_kernel.cpp new file mode 100644 index 00000000000..3b1d24fcbf8 --- /dev/null +++ b/src/plugins/intel_cpu/src/nodes/kernels/x64/mlp_kernel.cpp @@ -0,0 +1,385 @@ +// Copyright (C) 2022 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include "mlp_kernel.hpp" +#include "emitters/plugin/x64/jit_dnnl_emitters.hpp" +#include "mlp_utils.hpp" + +using namespace dnnl::impl; +using namespace dnnl::impl::utils; +using namespace dnnl::impl::cpu::x64; + +using TileConfig = ov::Extensions::Cpu::TileConfig; +using TileConfiger = ov::Extensions::Cpu::TileConfiger; + +namespace ov { +namespace intel_cpu { + +void MKernel::generate() { + Xbyak::Reg64 reg_A_addr = abi_param2; + Xbyak::Reg64 reg_A_stride = abi_param3; + Xbyak::Reg64 reg_B_addr = abi_param4; + Xbyak::Reg64 reg_C_addr = abi_param5; + Xbyak::Reg64 reg_C_stride = abi_param6; + + Xbyak::Reg64 reg_ktiles = rax; + Xbyak::Reg64 reg_B_stride = r10; + Xbyak::Reg64 reg_A1_addr = r11; + Xbyak::Reg64 reg_prefetch = r12; + + Xbyak::Tmm tmmC00 = tmm0; + Xbyak::Tmm tmmC01 = tmm1; + Xbyak::Tmm tmmC10 = tmm2; + Xbyak::Tmm tmmC11 = tmm3; + Xbyak::Tmm tmmA0 = tmm4; + Xbyak::Tmm tmmA1 = tmm5; + Xbyak::Tmm tmmB0 = tmm6; + Xbyak::Tmm tmmB1 = tmm7; + + auto num_PFB = m_prefetch_Blines; + int cur_PFB = 0; + + Xbyak::Label loop_over_ktiles; + Xbyak::Label skip_load; + + push(reg_prefetch); + { + auto reg_tmp = reg_B_stride; + tilezero(tmmC00); + tilezero(tmmC01); + tilezero(tmmC10); + tilezero(tmmC11); + + mov(reg_A_addr, ptr[abi_param1 + offsetof(call_args, pA)]); + mov(reg_A_stride, ptr[abi_param1 + offsetof(call_args, strideA)]); + mov(reg_B_addr, ptr[abi_param1 + offsetof(call_args, pB)]); + mov(reg_C_addr, ptr[abi_param1 + offsetof(call_args, pC)]); + mov(reg_C_stride, ptr[abi_param1 + offsetof(call_args, strideC)]); + mov(reg_prefetch, ptr[abi_param1 + offsetof(call_args, prefetch)]); + mov(reg_ktiles, ptr[abi_param1 + offsetof(call_args, k_tiles)]); + + lea(reg_A1_addr, ptr[reg_A_addr + reg_A_stride * 8]); + lea(reg_A1_addr, ptr[reg_A1_addr + reg_A_stride * 8]); + + // reg_A1_addr = reg_A_addr if M <= 16 (to avoid tileloadd segmentfault) + mov(reg_tmp, ptr[abi_param1 + offsetof(call_args, M)]); + cmp(reg_tmp, 16); + cmovle(reg_A1_addr, reg_A_addr); + + mov(reg_tmp, ptr[abi_param1 + offsetof(call_args, do_accumulation)]); + and_(reg_tmp, 1); + jz(skip_load); + { + auto reg_C1_addr = reg_tmp; + tileloadd(tmmC00, ptr[reg_C_addr + reg_C_stride]); + tileloadd(tmmC01, ptr[reg_C_addr + reg_C_stride + 64]); + lea(reg_C1_addr, ptr[reg_C_addr + reg_C_stride * 8]); + lea(reg_C1_addr, ptr[reg_C1_addr + reg_C_stride * 8]); + tileloadd(tmmC10, ptr[reg_C1_addr + reg_C_stride]); + tileloadd(tmmC11, ptr[reg_C1_addr + reg_C_stride + 64]); + } + L(skip_load); + } + + mov(reg_B_stride, 64); + + auto const_A_steps = 64; + + align(64, false); + L(loop_over_ktiles); + { + // B: 1x2 tiles + // A : 2x1 tiles C: 2x2 tiles + tileloadd(tmmA0, ptr[reg_A_addr + reg_A_stride]); + tileloadd(tmmB0, ptr[reg_B_addr + reg_B_stride]); + lea(reg_B_addr, ptr[reg_B_addr + 1024]); + + tdpbf16ps(tmmC00, tmmA0, tmmB0); + if (cur_PFB < num_PFB) { + prefetcht2(ptr[reg_prefetch + cur_PFB * 64]); + cur_PFB++; + } + + tileloadd(tmmA1, ptr[reg_A1_addr + reg_A_stride]); + tdpbf16ps(tmmC10, tmmA1, tmmB0); + if (cur_PFB < num_PFB) { + prefetcht2(ptr[reg_prefetch + cur_PFB * 64]); + cur_PFB++; + } + + tileloadd(tmmB1, ptr[reg_B_addr + reg_B_stride]); + tdpbf16ps(tmmC01, tmmA0, tmmB1); + if (cur_PFB < num_PFB) { + prefetcht2(ptr[reg_prefetch + cur_PFB * 64]); + cur_PFB++; + } + + tdpbf16ps(tmmC11, tmmA1, tmmB1); + if (cur_PFB < num_PFB) { + for (int pi = cur_PFB; pi < num_PFB; pi++) { + prefetcht2(ptr[reg_prefetch + pi * 64]); + } + } + + lea(reg_prefetch, ptr[reg_prefetch + 64 * num_PFB]); + lea(reg_A_addr, ptr[reg_A_addr + const_A_steps]); + lea(reg_A1_addr, ptr[reg_A1_addr + const_A_steps]); + lea(reg_B_addr, ptr[reg_B_addr + 1024]); + } + dec(reg_ktiles); + jnz(loop_over_ktiles, T_NEAR); + + tilestored(ptr[reg_C_addr + reg_C_stride], tmmC00); + tilestored(ptr[reg_C_addr + reg_C_stride + 64], tmmC01); + lea(reg_C_addr, ptr[reg_C_addr + reg_C_stride * 8]); + lea(reg_C_addr, ptr[reg_C_addr + reg_C_stride * 8]); + tilestored(ptr[reg_C_addr + reg_C_stride], tmmC10); + tilestored(ptr[reg_C_addr + reg_C_stride + 64], tmmC11); + + pop(reg_prefetch); + ret(); +} + +void MKernel::tile_config_M(TileConfig& tile_cfg, int M) { + auto rows0 = 16; + auto rows1 = 16; + if (M < 32) { + // kernel is for processing Mtails + if (M > 16) { + rows0 = 16; + rows1 = M - 16; + } else { + // both A0 & A1 load from same memory, to avoid code-regeneration + rows0 = rows1 = M; + } + } + tile_cfg.reset(1, + 0, + { + {rows0, 64}, // C00:0 + {rows0, 64}, // C01:1 + {rows1, 64}, // C10:2 + {rows1, 64}, // C11:3 + {rows0, 64}, // A0:4 + {rows1, 64}, // A1:5 + {16, 64}, // B0:6 + {16, 64}, // B1:7 + }); +} + +template +void MKernel::repackB(ov::bfloat16* dst, T* src, int N_stride, int N, int K) { + if (N == 16 && K == 32 && std::is_same::value) { + // SIMD optimized version + ov::Extensions::Cpu::XARCH::llm_mlp_transpose_epi32_16x16(dst, src, N_stride * sizeof(T)); + return; + } + + assert(K <= 32); + assert(N <= 16); + int k = 0; + ov::bfloat16 bf16zero(0.0f); + for (; k < 32; k += 2) { + int n = 0; + bool is_k0_valid = (k) < K; + bool is_k1_valid = (k + 1) < K; + auto* psrc = src + k; + for (; n < 16 && n < N; n++, psrc += N_stride) { + *dst++ = is_k0_valid ? ov::bfloat16(psrc[0]) : bf16zero; + *dst++ = is_k1_valid ? ov::bfloat16(psrc[1]) : bf16zero; + } + for (; n < 16; n++) { + *dst++ = 0; + *dst++ = 0; + } + } +} + +template +void MKernel::prepareB(PlainTensor& ret, T* p_weight, int stride, int N, int K) { + OPENVINO_ASSERT((N % 32) == 0); + OPENVINO_ASSERT((K % 32) == 0); + // weight matrix is in unit of [N/32, Kx32] + ret.resize({static_cast(N / 32), static_cast(K * 32)}); + + auto N_stride = stride / sizeof(T); + for (int n = 0, blkn = 0; n < N; n += 32, blkn++) { + for (int k = 0, blkk = 0; k < K; k += 32, blkk++) { + // two adjacent 32x16 (512) block of weight: dst0 & dst1 + auto* dst0 = ret.ptr(blkn, blkk * 1024); + auto* dst1 = dst0 + 16 * 32; + auto valid_k = (K - k) < 32 ? (K - k) : 32; + + auto* src0 = p_weight + n * N_stride + k; + auto valid_n0 = (N - n) < 16 ? (N - n) : 16; + repackB(dst0, src0, N_stride, valid_n0, valid_k); + + auto* src1 = p_weight + (n + 16) * N_stride + k; + auto valid_n1 = (N - (n + 16)) < 16 ? (N - (n + 16)) : 16; + repackB(dst1, src1, N_stride, valid_n1, valid_k); + } + } +} + +template void MKernel::prepareB(PlainTensor& ret, ov::bfloat16* p_weight, int stride, int N, int K); + +// run L2 cache blocking kernel with size: +// [BM, BK]*[BK, BN] => [BM, BN] +// +// prefetch of A can be done inside of this level of kernel +// since it's done in unit of 32-rows +// but prefetch of next B must be specified by caller. +// +void MKernel::run(int M, // actual M + uint8_t* pA, + int strideA, // A [M, K] + PlainTensor& repacked_B, // B [N/32, K*32] ov::bfloat16 + uint8_t* pC, + int strideC, // C [M, N] + uint8_t* prefetch_B, // prefetch B + bool do_accumulation) { + call_args args; + // number of blocks in N dimension (in unit of 32 columns) + auto num_blkN = static_cast(repacked_B.size(0)); + auto K = repacked_B.size(1) / 32; + auto* pB = repacked_B.ptr(); + auto strideB = repacked_B.stride_bytes(0); + + args.do_accumulation = do_accumulation; + args.k_tiles = K / 32; + args.strideA = strideA; + args.strideC = strideC; + args.prefetch = prefetch_B; + assert((K % 32) == 0); + + auto prefetch_step = m_prefetch_Blines * 64 * args.k_tiles; + + // if (BM != m_BM_hint) it only effect prefetch of B which is not vital to function + for (int m = 0; m < M; m += 32, pA += 32 * strideA, pC += 32 * strideC) { + args.pB = pB; + args.M = std::min(M - m, 32); + args.pA = pA; + for (int ni = 0; ni < num_blkN; ni++, args.pB += strideB, args.prefetch += prefetch_step) { + args.pC = pC + ni * 32 * sizeof(float); + (*this)(&args); + } + } +} + +void GateUpCombine::generate() { + Xbyak::Label loop_begin; + + Xbyak::Reg64 src = abi_param1; + Xbyak::Reg64 dst = abi_param2; + Xbyak::Reg64 prefetch_dst = abi_param3; + Xbyak::Reg64 BN = abi_param4; + + Xbyak::Reg64 loop_i = rax; + const auto zmm_gate = zmm5; + const auto zmm_silu = zmm6; + const auto zmm_up = zmm0; + const auto ymm_dst = ymm5; + + // when save_state is false, push/pop will not be generated. + auto injector = std::make_shared>( + this, + m_act_alg, + 1.f, + 1.0f, + 1.f, + false, // save_state, state will be saved in our function + Xbyak::Reg64(Xbyak::Operand::R10), // p_table + Xbyak::Opmask(1), // k_mask + true, // is_fwd + false, // use_dst + false, // preserve_vmm + false); // preserve_p_table + + xor_(loop_i, loop_i); + injector->load_table_addr(); + + shr(BN, 1); // BN = BN/2; + align(64); + L(loop_begin); + { + vmovups(zmm_gate, ptr[src + loop_i * 8]); + // silu will internally use zmm0~zmm3, gelu will use ~zmm4 + vmovups(zmm_silu, zmm_gate); + injector->compute_vector(zmm_silu.getIdx()); + vmovups(zmm_up, ptr[src + loop_i * 8 + 16 * 4]); + vmulps(zmm_up, zmm_up, zmm_silu); + vcvtneps2bf16(ymm_dst, zmm_up); + prefetchwt1(ptr[prefetch_dst + loop_i * 2]); + vmovdqu(ptr[dst + loop_i * 2], ymm_dst); + } + add(loop_i, 16); + cmp(loop_i, BN); + jl(loop_begin, T_NEAR); + + ret(); + + injector->prepare_table(); +} + +void ReduceAdd2bh::generate() { + if (m_do_reduce2) { + Xbyak::Reg64 src0 = abi_param1; + Xbyak::Reg64 src1 = abi_param2; + Xbyak::Reg64 dst = abi_param3; + Xbyak::Reg64 prefetch_dst = abi_param4; + Xbyak::Reg64 BN = abi_param5; + Xbyak::Reg64 loop_i = rax; + + Xbyak::Label loop_begin; + + xor_(loop_i, loop_i); + + align(64, false); + L(loop_begin); + { + vmovups(zmm0, ptr[src0 + loop_i * 4]); + vmovups(zmm1, ptr[src1 + loop_i * 4]); + vmovups(zmm2, ptr[src0 + loop_i * 4 + 16 * 4]); + vmovups(zmm3, ptr[src1 + loop_i * 4 + 16 * 4]); + vaddps(zmm0, zmm0, zmm1); + vaddps(zmm2, zmm2, zmm3); + vcvtne2ps2bf16(zmm4, zmm2, zmm0); + prefetchwt1(ptr[prefetch_dst + loop_i * 2]); + vmovups(ptr[dst + loop_i * 2], zmm4); + } + add(loop_i, 32); + cmp(loop_i, BN); + jl(loop_begin, T_NEAR); + + ret(); + } else { + Xbyak::Reg64 src0 = abi_param1; + Xbyak::Reg64 dst = abi_param2; + Xbyak::Reg64 prefetch_dst = abi_param3; + Xbyak::Reg64 BN = abi_param4; + Xbyak::Reg64 loop_i = rax; + + Xbyak::Label loop_begin; + + xor_(loop_i, loop_i); + + align(64, false); + L(loop_begin); + { + vmovups(zmm0, ptr[src0 + loop_i * 4]); + vmovups(zmm2, ptr[src0 + loop_i * 4 + 16 * 4]); + vcvtne2ps2bf16(zmm4, zmm2, zmm0); + prefetchwt1(ptr[prefetch_dst + loop_i * 2]); + vmovups(ptr[dst + loop_i * 2], zmm4); + } + add(loop_i, 32); + cmp(loop_i, BN); + jl(loop_begin, T_NEAR); + + ret(); + } +} + +} // namespace intel_cpu +} // namespace ov diff --git a/src/plugins/intel_cpu/src/nodes/kernels/x64/mlp_kernel.hpp b/src/plugins/intel_cpu/src/nodes/kernels/x64/mlp_kernel.hpp new file mode 100644 index 00000000000..5c969fb7837 --- /dev/null +++ b/src/plugins/intel_cpu/src/nodes/kernels/x64/mlp_kernel.hpp @@ -0,0 +1,264 @@ +// Copyright (C) 2022 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#pragma once + +#include "cpu/x64/cpu_isa_traits.hpp" +#include "cpu/x64/jit_generator.hpp" +#include "../scaled_attn/executor_pa_common.hpp" +#include "utils/plain_tensor.hpp" + +namespace ov { +namespace intel_cpu { + +class AutoTileConfiger { +public: + AutoTileConfiger() {} + ~AutoTileConfiger() { + do_config(nullptr); + } + void do_config(void* cfg) { + static ov::Extensions::Cpu::TileConfiger configer; + if (cfg != last_cfg) { + configer(cfg); + last_cfg = cfg; + } + } + +private: + void* last_cfg = nullptr; +}; + +class MKernel : public dnnl::impl::cpu::x64::jit_generator { +public: + DECLARE_CPU_JIT_AUX_FUNCTIONS(MKernel) + + int m_prefetch_Blines; + + MKernel(int M_hint = 256) : jit_generator("MKernel") { + setup(M_hint); + } + + void generate() override; + + // M_hint is only a hint for prefetching, set to 0 to avoid prefetch + void setup(int M_hint = 0) { + if (M_hint == 0) { + m_prefetch_Blines = 0; + } else { + m_prefetch_Blines = 32768 * sizeof(ov::bfloat16) / 64 / M_hint; + } + + create_kernel(); + } + + // M can change w/o code-regeneration + // with the help of : + // - m_BM_hint controls dynamic behaviour of the kernel + // - tile config controls behaviour of tileload & TMUL + void tile_config_M(ov::Extensions::Cpu::TileConfig& tile_cfg, int M); + + // row data is in layout [N, K], maybe smaller than [32, 16] + template + void repackB(ov::bfloat16* dst, T* src, int N_stride, int N, int K); + + // weight is supposed to be of shape[N, K], stride in unit of bytes + template + void prepareB(PlainTensor& ret, T* p_weight, int stride, int N, int K); + + // to save push/pop: do not use `abi_save_gpr_regs` + uint8_t* prefetch_next_A_addr; + + struct call_args { + const uint8_t* pA; // bfloat16 + int64_t strideA; // in bytes + const uint8_t* pB; // bfloat16 + const uint8_t* pC; // float32 + int64_t strideC; // in bytes + const uint8_t* prefetch; + int64_t k_tiles; // K / 32 + int64_t do_accumulation; + int64_t M; + }; + + // run L2 cache blocking kernel with size: + // [BM, BK]*[BK, BN] => [BM, BN] + // + // prefetch of A can be done inside of this level of kernel + // since it's done in unit of 32-rows + // but prefetch of next B must be specified by caller. + // + void run(int M, // actual M + uint8_t* pA, + int strideA, // A [M, K] + PlainTensor& repacked_B, // B [N/32, K*32] ov::bfloat16 + uint8_t* pC, + int strideC, // C [M, N] + uint8_t* prefetch_B, // prefetch B + bool do_accumulation); +}; + + +struct Work { + std::vector weights; // ov::bfloat16 weights for current thread + + std::shared_ptr sync_flag; + int n0 = 0; + int n1 = 0; + int k0 = 0; + int k1 = 0; + int BN = 0; + int blk_K_size = 0; + int output_id; + ov::bfloat16* p_raw_weights; + operator bool() { + return BN > 0; + } + + MKernel& get_MKernel() { + constexpr int BM = 256; + static MKernel jit_amx0(BM); + return jit_amx0; + } + + // input : weight [N, K], setup repacks range of N [n_start, n_end) + template + void setup(T* p_weight, int stride) { + auto& mkernel = get_MKernel(); + auto num_blk_K = (k1 - k0) / blk_K_size; + auto* pw = p_weight + n0 * stride / sizeof(T) + k0; + + weights.resize(num_blk_K); + for (int k = 0; k < num_blk_K; k++) { + mkernel.prepareB(weights[k], pw + k * blk_K_size, stride, BN, blk_K_size); + } + + for (int Mtails = 0; Mtails < 32; Mtails++) { + mkernel.tile_config_M(m_tcfg[Mtails], Mtails == 0 ? 32 : Mtails); + } + } + + ov::Extensions::Cpu::TileConfig m_tcfg[32]; + AutoTileConfiger m_tile_configer; + + size_t get_C_size(int M) { + auto Mtails = M % 32; + auto Mbody = M - Mtails; + auto C_M = Mbody + (Mtails ? 32 : 0); + return C_M * BN; + } + + void run(int M, uint8_t* pA, int strideA, PlainTensor& C) { + auto& mkernel = get_MKernel(); + + int num_blk_K = (k1 - k0) / blk_K_size; + + auto Mtails = M % 32; + auto Mbody = M - Mtails; + + auto C_M = Mbody + (Mtails ? 32 : 0); + C.resize({static_cast(C_M), static_cast(BN)}); + auto pC = reinterpret_cast(C.ptr_v()); + + pA += k0 * sizeof(ov::bfloat16); + bool do_accumulation = false; + + for (int ki = 0; ki < num_blk_K; ki++) { + PlainTensor& blockB = weights[ki]; + PlainTensor& blockB1 = weights[(ki + 1) < num_blk_K ? (ki + 1) : ki]; + if (Mbody) { + m_tile_configer.do_config(&m_tcfg[0]); + mkernel.run(Mbody, + pA + ki * blk_K_size * sizeof(ov::bfloat16), + strideA, + blockB, + pC, + C.stride_bytes(0), + reinterpret_cast(blockB1.ptr_v()), + do_accumulation); + } + + if (Mtails) { + m_tile_configer.do_config(&m_tcfg[Mtails]); + mkernel.run(Mtails, + pA + ki * blk_K_size * sizeof(ov::bfloat16) + Mbody * strideA, + strideA, + blockB, + pC + Mbody * C.stride_bytes(0), + C.stride_bytes(0), + reinterpret_cast(blockB1.ptr_v()), + do_accumulation); + } + do_accumulation = true; + } + m_tile_configer.do_config(nullptr); + } +}; + +// combine gate_proj & up_proj using activation algo, then convert to bf16 +// ConvertFP32toBF16(act_fn(gate) * up) +class GateUpCombine : public dnnl::impl::cpu::x64::jit_generator { +public: + DECLARE_CPU_JIT_AUX_FUNCTIONS(GateUpCombine) + + const dnnl_alg_kind_t m_act_alg; + GateUpCombine(dnnl_alg_kind_t act_alg) : jit_generator(jit_name()), m_act_alg(act_alg) { + create_kernel(); + } + + void generate() override; + + void call(float* src, size_t src_stride, ov::bfloat16 * dst, size_t dst_stride, int num_rows, int num_cols) { + for (int m = 0; m < num_rows; m++, src += src_stride, dst += dst_stride) { + auto* prefetch_dst = (m + 1 < num_rows) ? (dst + dst_stride) : (dst); + + // gate_proj & up_proj are interleaved in unit of 16 elements into gateup + // + // for(int i = 0; i < N; i += 32) { + // for(int k = 0; k < 16; k++) + // gate = src[i+k] + // up = src[i+k+16] + // *dst++ = ConvertFP32toBF16(act_fn(gate) * up_proj) + // } + // } + // + (*this)(src, dst, prefetch_dst, num_cols); + } + } +}; + +class ReduceAdd2bh : public dnnl::impl::cpu::x64::jit_generator { +public: + DECLARE_CPU_JIT_AUX_FUNCTIONS(ReduceAdd2bh) + + bool m_do_reduce2; + ReduceAdd2bh(bool do_reduce2) : jit_generator(jit_name()), m_do_reduce2(do_reduce2) { + create_kernel(); + } + + void generate() override; + + // add two float input eltwise and convert to bf16 : ConvertFP32toBF16(src0 + src1) + void call(float * src0, float * src1, size_t src_stride, ov::bfloat16 * dst, size_t dst_stride, int num_rows, int num_cols) { + for (int m = 0; m < num_rows; m++, src0 += src_stride, src1 += src_stride, dst += dst_stride) { + // the prefetch distance is increased to ensure by the time store happens + // prefetch has done and no HW prefetcher is triggered + auto* prefetch_dst = (m + 2 < num_rows) ? (dst + 2 * dst_stride) : (dst); + (*this)(src0, src1, dst, prefetch_dst, num_cols); + } + } + + // convert tensor to bf16: ConvertFP32toBF16(src0) + void call(float * src0, size_t src_stride, ov::bfloat16 * dst, size_t dst_stride, int num_rows, int num_cols) { + for (int m = 0; m < num_rows; m++, src0 += src_stride, dst += dst_stride) { + // the prefetch distance is increased to ensure by the time store happens + // prefetch has done and no HW prefetcher is triggered + auto* prefetch_dst = (m + 2 < num_rows) ? (dst + 2 * dst_stride) : (dst); + (*this)(src0, dst, prefetch_dst, num_cols); + } + } +}; + +} // namespace intel_cpu +} // namespace ov diff --git a/src/plugins/intel_cpu/src/nodes/kernels/x64/mlp_utils.cpp b/src/plugins/intel_cpu/src/nodes/kernels/x64/mlp_utils.cpp new file mode 100644 index 00000000000..996561a04af --- /dev/null +++ b/src/plugins/intel_cpu/src/nodes/kernels/x64/mlp_utils.cpp @@ -0,0 +1,28 @@ +// Copyright (C) 2018-2024 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include "mlp_utils.hpp" + + +#include +#if defined(HAVE_AVX512F) +#include +#endif +#include "../scaled_attn/transpose_kernel.hpp" + +namespace ov { +namespace Extensions { +namespace Cpu { +namespace XARCH { + +void llm_mlp_transpose_epi32_16x16(void* dst, void* src, int stride) { + transpose_16x16_kernel(reinterpret_cast(dst), reinterpret_cast(src), 16, stride/sizeof(uint32_t)); +} + +} // namespace XARCH +} // namespace Cpu +} // namespace Extensions +} // namespace ov + + diff --git a/src/plugins/intel_cpu/src/nodes/kernels/x64/mlp_utils.hpp b/src/plugins/intel_cpu/src/nodes/kernels/x64/mlp_utils.hpp new file mode 100644 index 00000000000..075717971f4 --- /dev/null +++ b/src/plugins/intel_cpu/src/nodes/kernels/x64/mlp_utils.hpp @@ -0,0 +1,19 @@ +// Copyright (C) 2018-2024 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include +#include +#include + +namespace ov { +namespace Extensions { +namespace Cpu { +namespace XARCH { + +void llm_mlp_transpose_epi32_16x16(void* dst, void* src, int stride); + +} // namespace XARCH +} // namespace Cpu +} // namespace Extensions +} // namespace ov diff --git a/src/plugins/intel_cpu/src/nodes/llm_mlp.cpp b/src/plugins/intel_cpu/src/nodes/llm_mlp.cpp new file mode 100644 index 00000000000..dc63bdf740e --- /dev/null +++ b/src/plugins/intel_cpu/src/nodes/llm_mlp.cpp @@ -0,0 +1,339 @@ +// Copyright (C) 2018-2024 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include "llm_mlp.h" + +#include +#include + +#include "common/bfloat16.hpp" +#include "common/cpu_memcpy.h" +#include "cpu/x64/cpu_isa_traits.hpp" +#include "cpu/x64/jit_generator.hpp" +#include "shape_inference/shape_inference_internal_dyn.hpp" +#include "utils/plain_tensor.hpp" + +#if defined(OPENVINO_ARCH_X86_64) +#include "kernels/x64/mlp_kernel.hpp" +#endif + +namespace ov { +namespace intel_cpu { +namespace node { + +#if defined(OPENVINO_ARCH_X86_64) + +class Linear { +public: + std::vector works; + + int used_nthr = 0; + bool do_splitK = false; + + Linear() {} + + // weight [N, K] + // Gate & Up are interleaved in N dimension: 16-gate / 16-up + // and post-ops will compute silu(gate)*up in unit of 16 elements + // and store out as bfloat16. + template + void setup(T* p_weight, int stride, int N, int K, bool _do_splitK = false) { + const int blk_K_size = 256; + // prepare weights, split N among threads + // in unit of 32 + OPENVINO_ASSERT((N % 32) == 0); + OPENVINO_ASSERT((K % blk_K_size) == 0); + auto nthr = parallel_get_max_threads(); + auto num_blk_N = N / 32; + auto num_blk_K = K / blk_K_size; + works.resize(nthr); + + do_splitK = _do_splitK; + auto K_splits = do_splitK ? 2 : 1; + // split task on more cores is better on TBB + auto valid_nthr = nthr / K_splits; + auto blkN_per_thread = (num_blk_N) / valid_nthr; + auto blkN_leftover = num_blk_N - (blkN_per_thread * valid_nthr); + auto start_blkN = 0; + used_nthr = 0; + auto blkK_per_thread = (num_blk_K + K_splits - 1) / K_splits; + + for (int ithr = 0; ithr < nthr; ithr += K_splits) { + auto blkN = std::min(num_blk_N - start_blkN, blkN_per_thread); + if (blkN_leftover > 0) { + blkN_leftover--; + blkN++; + } + if (blkN) { + auto shared_atomic = std::make_shared(0); + auto start_blkK = 0; + for (int ik = 0; ik < K_splits; ik++) { + auto blk_K = std::min(num_blk_K - start_blkK, blkK_per_thread); + + auto& work = works[ithr + ik]; + + work.sync_flag = shared_atomic; + work.blk_K_size = blk_K_size; + + work.n0 = (start_blkN)*32; + work.n1 = (start_blkN + blkN) * 32; + work.BN = blkN * 32; + work.k0 = start_blkK * blk_K_size; + work.k1 = (start_blkK + blk_K) * blk_K_size; + + start_blkK += blk_K; + used_nthr++; + } + } + + start_blkN += blkN; + } + + DEBUG_LOG("Linear N,K=", N, ",", K, " used_nthr=", used_nthr, " do_splitK=", do_splitK); + + ov::parallel_nt_static(0, [&](const size_t ithr, const size_t nthr) { + auto& work = works[ithr]; + if (work) { + work.setup(p_weight, stride); + } + }); + DEBUG_LOG(" setup is done. weight @ ", static_cast(p_weight)); + } + + void run(uint8_t* pA, int strideA, int M, ov::bfloat16* dstC, int strideC, std::vector& m_tempC) { + static ReduceAdd2bh jit_reduce2bh_1(false); + static ReduceAdd2bh jit_reduce2bh_2(true); + + ov::parallel_nt_static(0, [&](const size_t ithr, const size_t nthr) { + auto& work = works[ithr]; + auto& workC = m_tempC[ithr]; + if (work) { + work.run(M, pA, strideA, workC); + + if (do_splitK) { + auto sync_id = work.sync_flag->fetch_add(1); + // (0,1) (2,3) + if (sync_id & 1) { + auto peer_ithr = (ithr & 1) ? (ithr - 1) : (ithr + 1); + auto& peerC = m_tempC[peer_ithr]; + // the other one has finished, we can do the reduce sum + jit_reduce2bh_2.call(workC.ptr(), peerC.ptr(), workC.stride(0), + dstC + work.n0, strideC / sizeof(*dstC), + M, work.BN); + } + } else { + jit_reduce2bh_2.call(workC.ptr(), workC.stride(0), + dstC + work.n0, strideC / sizeof(*dstC), + M, work.BN); + } + } + }); + } + + // gate & up are interleaved: 16 gates + 16 up + void runGateUp(uint8_t* pA, int strideA, int M, + ov::bfloat16* dstC, int strideC, + const LLMMLPNode::Config& config, + std::vector& m_tempC) { + static GateUpCombine jit_gateup_silu(dnnl_eltwise_swish); + static GateUpCombine jit_gateup_gelu(dnnl_eltwise_gelu_tanh); + + GateUpCombine* jit_gateup; + if (config.act == LLMMLPNode::ACT_FN::GELU) + jit_gateup = &jit_gateup_gelu; + else if (config.act == LLMMLPNode::ACT_FN::SILU) + jit_gateup = &jit_gateup_silu; + else + OPENVINO_THROW("unsupported act in GateUpCombine"); + + ov::parallel_nt_static(0, [&](const size_t ithr, const size_t nthr) { + auto& work = works[ithr]; + auto& workC = m_tempC[ithr]; + if (work.BN > 0) { + work.run(M, pA, strideA, workC); + // K reduce is done, results of [M, BN] sub-block is ready in L2. + // combine Gate & Up + jit_gateup->call(workC.ptr(), workC.stride(0), + dstC + (work.n0 / 2), strideC / sizeof(*dstC), + M, work.BN); + } + }); + } +}; + +struct LLMMLP::Impl { + const LLMMLPNode::Config m_config; + DnnlScratchPadPtr m_scrachPad; + MemoryPtr m_scratchMem; + + Linear gate_up; + Linear down; + int m_N; + int m_M = 0; + + // MLP is not supposed to run in parallel + PlainTensor m_actUp; + std::vector m_tempC; + + // [M, K] x [N, K] => [M, N] x [K, N] => [M, K] + // w_gate/w_up : [N, K] + // w_down : [K, N] + Impl(PlainTensor w_gate, PlainTensor w_up, PlainTensor w_down, const LLMMLPNode::Config& config, DnnlScratchPadPtr scrachPad) + : m_config(config), m_scrachPad(scrachPad) { + // [N, K] [N, K] interleave (16-16-...) into [2*N, K] + auto K = w_gate.size(1); + auto N = w_gate.size(0); + static PlainTensor w_gate_up; + w_gate_up.resize({static_cast(2 * N), static_cast(K)}); + for (size_t n = 0; n < N; n += 16) { + for (size_t i = 0; i < 16; i++) + memcpy(w_gate_up.ptr_v(2 * n + i, 0), w_gate.ptr_v(n + i, 0), K * sizeof(ov::bfloat16)); + for (size_t i = 0; i < 16; i++) + memcpy(w_gate_up.ptr_v(2 * n + 16 + i, 0), w_up.ptr_v(n + i, 0), K * sizeof(ov::bfloat16)); + } + gate_up.setup(w_gate_up.ptr(), w_gate_up.stride_bytes(0), N * 2, K); + down.setup(w_down.ptr(), w_down.stride_bytes(0), K, N, true); + + m_tempC.resize(parallel_get_max_threads()); + m_N = N; + } + + void setM(int M) { + if (m_M < M) { + size_t total_scratch_size = M * m_N * sizeof(ov::bfloat16); + std::vector scratch_offsets; + std::vector scratch_C_sizes; + for (size_t ithr = 0; ithr < m_tempC.size(); ithr++) { + scratch_offsets.push_back(total_scratch_size); + auto max_C_size = std::max(gate_up.works[ithr].get_C_size(M), down.works[ithr].get_C_size(M)); + scratch_C_sizes.push_back(max_C_size); + total_scratch_size += max_C_size * sizeof(float); + } + + auto newMemDesc = std::make_shared(ov::element::u8, Shape{total_scratch_size}); + m_scratchMem = m_scrachPad->createScratchPadMem(newMemDesc); + + auto* scratch_base = m_scratchMem->getDataAs(); + m_actUp.resize({static_cast(M), static_cast(m_N)}, reinterpret_cast(scratch_base)); + + for (size_t ithr = 0; ithr < m_tempC.size(); ithr++) { + m_tempC[ithr].resize({1, scratch_C_sizes[ithr]}, reinterpret_cast(scratch_base + scratch_offsets[ithr])); + } + m_M = M; + } + } + + void execute(LLMMLP* pnode) { + auto input = pnode->getSrcMemoryAtPort(0); + const auto& ishape = input->getStaticDims(); + uint8_t* pA = input->getDataAs(); + const auto& srcStrides = input->getDescWithType()->getStrides(); + + int strideA = srcStrides[srcStrides.size() - 2] * sizeof(ov::bfloat16); + int M = shape_size(ishape) / ishape[ishape.size() - 1]; + + auto output = pnode->getDstMemoryAtPort(0); + auto* dstC = output->getDataAs(); + const auto& dstStrides = output->getDescWithType()->getStrides(); + int strideC = dstStrides[dstStrides.size() - 2] * sizeof(ov::bfloat16); + + for (int m = 0; m < M;) { + int BM = std::min(M - m, 512); + setM(BM); + + gate_up.runGateUp(pA, strideA, BM, m_actUp.ptr(), m_actUp.stride_bytes(0), m_config, m_tempC); + down.run(reinterpret_cast(m_actUp.ptr()), m_actUp.stride_bytes(0), BM, dstC, strideC, m_tempC); + + m += BM; + pA += BM * strideA; + dstC += BM * strideC / sizeof(ov::bfloat16); + } + } +}; +#else +struct LLMMLP::Impl { + Impl(PlainTensor w_gate, PlainTensor w_up, PlainTensor w_down, const LLMMLPNode::Config& config, DnnlScratchPadPtr scrachPad) {} + void execute(LLMMLP* pnode) {} +}; +#endif + +LLMMLP::LLMMLP(const std::shared_ptr& op, const GraphContext::CPtr context) + : Node(op, context, NgraphShapeInferFactory(op, EMPTY_PORT_MASK)) { + std::string errorMessage; + if (!isSupportedOperation(op, errorMessage)) { + OPENVINO_THROW("CPU: " + errorMessage); + } + const auto node_mlp = std::dynamic_pointer_cast(op); + m_mlp_config = node_mlp->get_config(); +} + +void LLMMLP::initSupportedPrimitiveDescriptors() { + if (!supportedPrimitiveDescriptors.empty()) + return; + + auto rtPrecision = ov::element::bf16; + auto weightPrecision = ov::element::bf16; + + // initialize input ports + std::vector inPortConfigs; + inPortConfigs.emplace_back(LayoutType::ncsp, rtPrecision, getInputShapeAtPort(0), false, -1); // input + inPortConfigs.emplace_back(LayoutType::ncsp, weightPrecision, getInputShapeAtPort(1), false, -1); // gate + inPortConfigs.emplace_back(LayoutType::ncsp, weightPrecision, getInputShapeAtPort(2), false, -1); // up + inPortConfigs.emplace_back(LayoutType::ncsp, weightPrecision, getInputShapeAtPort(3), false, -1); // down + + // initialize output port + std::vector outPortConfigs; + outPortConfigs.emplace_back(LayoutType::ncsp, rtPrecision, getOutputShapeAtPort(0), false, -1); + + addSupportedPrimDesc(inPortConfigs, outPortConfigs, impl_desc_type::ref_any); +} + +void LLMMLP::prepareParams() { + if (!m_pimpl) { + m_pimpl = std::make_shared(getSrcMemoryAtPort(1), + getSrcMemoryAtPort(2), + getSrcMemoryAtPort(3), + m_mlp_config, + context->getScratchPad()); + } +} + +void LLMMLP::execute(dnnl::stream strm) { + MAYBE_UNUSED(strm); + m_pimpl->execute(this); +} + +bool LLMMLP::isSupportedOperation(const std::shared_ptr& op, std::string& errorMessage) noexcept { + try { + const auto node_mlp = std::dynamic_pointer_cast(op); + if (node_mlp) { + auto down_proj_w_pshape = op->input_value(1).get_partial_shape(); + if (!down_proj_w_pshape.is_static()) { + // return true to skip Fusion + errorMessage = "LLMMLPNode weight shape is not static"; + return false; + } + auto down_size = down_proj_w_pshape[0].get_length(); + auto up_size = down_proj_w_pshape[1].get_length(); + if ((down_size % 256) != 0) { + errorMessage = "LLMMLPNode down_proj size is not multiple of 256"; + return false; + } + if (up_size % 256) { + errorMessage = "LLMMLPNode up_proj size is not multiple of 256"; + return false; + } + } else { + errorMessage = "Only LLMMLPNode operation is supported"; + return false; + } + } catch (...) { + return false; + } + return true; +} + +} // namespace node +} // namespace intel_cpu +} // namespace ov diff --git a/src/plugins/intel_cpu/src/nodes/llm_mlp.h b/src/plugins/intel_cpu/src/nodes/llm_mlp.h new file mode 100644 index 00000000000..3cdf6a0f1f7 --- /dev/null +++ b/src/plugins/intel_cpu/src/nodes/llm_mlp.h @@ -0,0 +1,39 @@ +// Copyright (C) 2018-2024 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#pragma once + +#include "node.h" +#include "transformations/cpu_opset/x64/op/llm_mlp.hpp" +#include "transformations/cpu_opset/x64/op/qkv_proj.hpp" + +namespace ov { +namespace intel_cpu { +namespace node { + +class LLMMLP : public Node { +public: + LLMMLP(const std::shared_ptr& op, const GraphContext::CPtr context); + + void getSupportedDescriptors() override {} + bool created() const override { + return getType() == Type::LLMMLP; + } + void prepareParams() override; + void executeDynamicImpl(dnnl::stream strm) override { + execute(strm); + } + void initSupportedPrimitiveDescriptors() override; + void execute(dnnl::stream strm) override; + static bool isSupportedOperation(const std::shared_ptr& op, std::string& errorMessage) noexcept; + +private: + struct Impl; + LLMMLPNode::Config m_mlp_config; + std::shared_ptr m_pimpl; +}; + +} // namespace node +} // namespace intel_cpu +} // namespace ov diff --git a/src/plugins/intel_cpu/src/nodes/qkv_proj.cpp b/src/plugins/intel_cpu/src/nodes/qkv_proj.cpp new file mode 100644 index 00000000000..74cfecf7aee --- /dev/null +++ b/src/plugins/intel_cpu/src/nodes/qkv_proj.cpp @@ -0,0 +1,291 @@ +// Copyright (C) 2018-2024 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include "qkv_proj.h" + +#include +#include + +#include "common/bfloat16.hpp" +#include "common/cpu_memcpy.h" +#include "cpu/x64/cpu_isa_traits.hpp" +#include "cpu/x64/jit_generator.hpp" +#include "shape_inference/shape_inference_internal_dyn.hpp" +#include "utils/plain_tensor.hpp" + +namespace ov { +namespace intel_cpu { +namespace node { + +#if defined(OPENVINO_ARCH_X86_64) +static std::vector allocate_workers(const std::vector& grouped_works, int n_workers) { + auto n_groups = grouped_works.size(); + // allocate 1 worker for each group + std::vector g_workers(n_groups, 1); + auto left_workers = n_workers - n_groups; + while (left_workers > 0) { + // which group is working hardest? + float hardest_works = 0; + size_t hardest_group = 0; + for (size_t g = 0; g < n_groups; g++) { + auto works = static_cast(grouped_works[g]) / g_workers[g]; + if (hardest_works < works) { + hardest_works = works; + hardest_group = g; + } + } + g_workers[hardest_group]++; + left_workers--; + } + + return g_workers; +} + +struct QKVProjection::Impl { + std::vector works; + std::vector m_tempC; + QKVProjection * m_node; + DnnlScratchPadPtr m_scrachPad; + MemoryPtr m_scratchMem; + int m_M; + + Impl(QKVProjection * pnode, DnnlScratchPadPtr scrachPad) : m_node(pnode), m_scrachPad(scrachPad) { + PlainTensor w0(pnode->getSrcMemoryAtPort(1)); + PlainTensor w1(pnode->getSrcMemoryAtPort(2)); + PlainTensor w2(pnode->getSrcMemoryAtPort(3)); + + const int blk_K_size = 256; + auto K = w0.size(1); + OPENVINO_ASSERT((K % blk_K_size) == 0); + auto nthr = parallel_get_max_threads(); + auto num_blk_K = K / blk_K_size; + int stride = K * sizeof(ov::bfloat16); + + works.resize(nthr); + + int cur_work_id = 0; + auto create_works = [&](ov::bfloat16* pw, int output_id, int N, int valid_nthr) { + // split task on more cores is better on TBB + OPENVINO_ASSERT((N % 32) == 0); + auto num_blk_N = N / 32; + auto blkN_per_thread = (num_blk_N) / valid_nthr; + auto blkN_leftover = num_blk_N - (blkN_per_thread * valid_nthr); + auto start_blkN = 0; + + for (int ithr = 0; ithr < valid_nthr; ithr++) { + auto blkN = std::min(num_blk_N - start_blkN, blkN_per_thread); + if (blkN_leftover > 0) { + blkN_leftover--; + blkN++; + } + if (blkN) { + auto& work = works[cur_work_id++]; + work.blk_K_size = blk_K_size; + work.n0 = (start_blkN)*32; + work.n1 = (start_blkN + blkN) * 32; + work.BN = blkN * 32; + work.k0 = 0; + work.k1 = blk_K_size * num_blk_K; + work.output_id = output_id; + work.p_raw_weights = pw; + } + start_blkN += blkN; + } + }; + auto proj_size0 = static_cast(w0.size(0)); + auto proj_size1 = static_cast(w1.size(0)); + auto proj_size2 = static_cast(w2.size(0)); + auto n_group_workers = allocate_workers({proj_size0, proj_size1, proj_size2}, nthr); + + create_works(w0.ptr(), 0, proj_size0, n_group_workers[0]); + create_works(w1.ptr(), 1, proj_size1, n_group_workers[1]); + create_works(w2.ptr(), 2, proj_size2, n_group_workers[2]); + + DEBUG_LOG("QKVProj hidden_size=", K, " proj_sizes=", + proj_size0, ",", proj_size1, ",", proj_size2, + " used_nthr=", cur_work_id); + ov::parallel_nt_static(0, [&](const size_t ithr, const size_t nthr) { + auto& work = works[ithr]; + if (work) { + work.setup(work.p_raw_weights, stride); + } + }); + + m_tempC.resize(nthr); + } + + void setM(int M) { + if (m_M < M) { + size_t total_scratch_size = 0; + std::vector scratch_offsets; + std::vector scratch_C_sizes; + for (size_t ithr = 0; ithr < m_tempC.size(); ithr++) { + scratch_offsets.push_back(total_scratch_size); + auto max_C_size = works[ithr].get_C_size(M); + scratch_C_sizes.push_back(max_C_size); + total_scratch_size += max_C_size * sizeof(float); + } + + auto newMemDesc = std::make_shared(ov::element::u8, Shape{total_scratch_size}); + m_scratchMem = m_scrachPad->createScratchPadMem(newMemDesc); + + auto* scratch_base = m_scratchMem->getDataAs(); + for (size_t ithr = 0; ithr < m_tempC.size(); ithr++) { + m_tempC[ithr].resize({1, scratch_C_sizes[ithr]}, reinterpret_cast(scratch_base + scratch_offsets[ithr])); + } + + m_M = M; + } + } + + void execute() { + static ReduceAdd2bh jit_2bh(false); + auto input = m_node->getSrcMemoryAtPort(0); + const auto& ishape = input->getStaticDims(); + uint8_t* pA = input->getDataAs(); + int M = shape_size(ishape) / ishape[ishape.size() - 1]; + auto* dst0 = m_node->getDstMemoryAtPort(0)->getDataAs(); + auto* dst1 = m_node->getDstMemoryAtPort(1)->getDataAs(); + auto* dst2 = m_node->getDstMemoryAtPort(2)->getDataAs(); + + const auto& srcStrides = input->getDescWithType()->getStrides(); + const auto& dstStrides0 = m_node->getDstMemoryAtPort(0)->getDescWithType()->getStrides(); + const auto& dstStrides1 = m_node->getDstMemoryAtPort(1)->getDescWithType()->getStrides(); + const auto& dstStrides2 = m_node->getDstMemoryAtPort(2)->getDescWithType()->getStrides(); + + int strideA = srcStrides[1] * sizeof(ov::bfloat16); + auto stride_0 = dstStrides0[1]; + auto stride_1 = dstStrides1[1]; + auto stride_2 = dstStrides2[1]; + + for (int m = 0; m < M;) { + int BM = std::min(M - m, 256); + + setM(BM); + + ov::parallel_nt_static(0, [&](const size_t ithr, const size_t nthr) { + auto& work = works[ithr]; + auto& C = m_tempC[ithr]; + if (work.BN > 0) { + work.run(BM, pA, strideA, C); + + // compress accumulation result into target + auto* src = C.ptr(); + auto stride_src = C.stride(0); + ov::bfloat16* dst = nullptr; + int stride_dst = 0; + if (work.output_id == 0) { + dst = dst0 + work.n0; + stride_dst = stride_0; + } + if (work.output_id == 1) { + dst = dst1 + work.n0; + stride_dst = stride_1; + } + if (work.output_id == 2) { + dst = dst2 + work.n0; + stride_dst = stride_2; + } + + for (int mi = 0; mi < BM; mi++, src += stride_src, dst += stride_dst) { + // the prefetch distance is increased to ensure by the time store happens + // prefetch has done and no HW prefetcher is triggered + auto* prefetch_dst = (mi + 2 < BM) ? (dst + 2 * stride_dst) : (dst); + jit_2bh(src, dst, prefetch_dst, work.BN); + } + } + }); + m += BM; + pA += BM * strideA; + dst0 += BM * stride_0 / sizeof(ov::bfloat16); + dst1 += BM * stride_1 / sizeof(ov::bfloat16); + dst2 += BM * stride_2 / sizeof(ov::bfloat16); + } + } +}; +#else +struct QKVProjection::Impl { + Impl(QKVProjection * pnode, DnnlScratchPadPtr scrachPad) {} + void execute() {} +}; +#endif + +void QKVProjection::prepareParams() { + if (!m_pimpl) { + m_pimpl = std::make_shared(this, context->getScratchPad()); + } +} + +void QKVProjection::execute(dnnl::stream strm) { + MAYBE_UNUSED(strm); + m_pimpl->execute(); +} + +QKVProjection::QKVProjection(const std::shared_ptr& op, const GraphContext::CPtr context) + : Node(op, context, NgraphShapeInferFactory(op, EMPTY_PORT_MASK)) { + std::string errorMessage; + if (!isSupportedOperation(op, errorMessage)) { + OPENVINO_THROW("CPU: " + errorMessage); + } +} + +void QKVProjection::initSupportedPrimitiveDescriptors() { + if (!supportedPrimitiveDescriptors.empty()) + return; + + auto rtPrecision = ov::element::bf16; + auto weightPrecision = ov::element::bf16; + + // initialize input ports + std::vector inPortConfigs; + inPortConfigs.emplace_back(LayoutType::ncsp, rtPrecision, getInputShapeAtPort(0), false, -1); // input + inPortConfigs.emplace_back(LayoutType::ncsp, weightPrecision, getInputShapeAtPort(1), false, -1); // gate + inPortConfigs.emplace_back(LayoutType::ncsp, weightPrecision, getInputShapeAtPort(2), false, -1); // up + inPortConfigs.emplace_back(LayoutType::ncsp, weightPrecision, getInputShapeAtPort(3), false, -1); // down + + // initialize output port + std::vector outPortConfigs; + outPortConfigs.emplace_back(LayoutType::ncsp, rtPrecision, getOutputShapeAtPort(0), false, -1); + outPortConfigs.emplace_back(LayoutType::ncsp, rtPrecision, getOutputShapeAtPort(1), false, -1); + outPortConfigs.emplace_back(LayoutType::ncsp, rtPrecision, getOutputShapeAtPort(2), false, -1); + + addSupportedPrimDesc(inPortConfigs, outPortConfigs, impl_desc_type::ref_any); +} + +bool QKVProjection::isSupportedOperation(const std::shared_ptr& op, std::string& errorMessage) noexcept { + try { + const auto node_qkv = std::dynamic_pointer_cast(op); + if (node_qkv) { + auto proj_pshape1 = op->input_value(1).get_shape(); + auto proj_pshape2 = op->input_value(2).get_shape(); + auto proj_pshape3 = op->input_value(3).get_shape(); + if ((proj_pshape1[1] % 256) != 0) { + errorMessage = "QKVProjection input channel size is not multiple of 256"; + return false; + } + if ((proj_pshape1[0] % 32) != 0) { + errorMessage = "QKVProjection 1st proj output channel size is not multiple of 32"; + return false; + } + if ((proj_pshape2[0] % 32) != 0) { + errorMessage = "QKVProjection 2nd proj output channel size is not multiple of 32"; + return false; + } + if ((proj_pshape3[0] % 32) != 0) { + errorMessage = "QKVProjection 3rd proj output channel size is not multiple of 32"; + return false; + } + } else { + errorMessage = "Only QKVProjection operation is supported"; + return false; + } + } catch (...) { + return false; + } + return true; +} + +} // namespace node +} // namespace intel_cpu +} // namespace ov diff --git a/src/plugins/intel_cpu/src/nodes/qkv_proj.h b/src/plugins/intel_cpu/src/nodes/qkv_proj.h new file mode 100644 index 00000000000..2d235a30505 --- /dev/null +++ b/src/plugins/intel_cpu/src/nodes/qkv_proj.h @@ -0,0 +1,41 @@ +// Copyright (C) 2018-2024 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#pragma once + +#include "node.h" +#include "transformations/cpu_opset/x64/op/qkv_proj.hpp" + +#if defined(OPENVINO_ARCH_X86_64) +#include "kernels/x64/mlp_kernel.hpp" +#endif + +namespace ov { +namespace intel_cpu { +namespace node { + +class QKVProjection : public Node { +public: + QKVProjection(const std::shared_ptr& op, const GraphContext::CPtr context); + + void getSupportedDescriptors() override {} + bool created() const override { + return getType() == Type::QKVProjection; + } + void prepareParams() override; + void executeDynamicImpl(dnnl::stream strm) override { + execute(strm); + } + void initSupportedPrimitiveDescriptors() override; + void execute(dnnl::stream strm) override; + static bool isSupportedOperation(const std::shared_ptr& op, std::string& errorMessage) noexcept; + +private: + struct Impl; + std::shared_ptr m_pimpl; +}; + +} // namespace node +} // namespace intel_cpu +} // namespace ov diff --git a/src/plugins/intel_cpu/src/nodes_factory.cpp b/src/plugins/intel_cpu/src/nodes_factory.cpp index 7833c60d071..f82f0adece1 100644 --- a/src/plugins/intel_cpu/src/nodes_factory.cpp +++ b/src/plugins/intel_cpu/src/nodes_factory.cpp @@ -43,6 +43,8 @@ #include "nodes/if.h" #include "nodes/input.h" #include "nodes/interaction.h" +#include "nodes/llm_mlp.h" +#include "nodes/qkv_proj.h" #include "nodes/interpolate.h" #include "nodes/inverse.hpp" #include "nodes/log_softmax.h" @@ -209,6 +211,8 @@ Node::NodesFactory::NodesFactory() : Factory("NodesFactory") { INTEL_CPU_NODE(FakeQuantize, Type::FakeQuantize); INTEL_CPU_NODE(GridSample, Type::GridSample); INTEL_CPU_NODE(Interaction, Type::Interaction); + INTEL_CPU_NODE(LLMMLP, Type::LLMMLP); + INTEL_CPU_NODE(QKVProjection, Type::QKVProjection); INTEL_CPU_NODE(MHA, Type::MHA); INTEL_CPU_NODE(ScaledDotProductAttention, Type::ScaledDotProductAttention); INTEL_CPU_NODE(PagedAttention, Type::PagedAttention); diff --git a/src/plugins/intel_cpu/src/transformations/cpu_opset/x64/op/llm_mlp.cpp b/src/plugins/intel_cpu/src/transformations/cpu_opset/x64/op/llm_mlp.cpp new file mode 100644 index 00000000000..2ad643d3ad4 --- /dev/null +++ b/src/plugins/intel_cpu/src/transformations/cpu_opset/x64/op/llm_mlp.cpp @@ -0,0 +1,59 @@ +// Copyright (C) 2018-2024 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include "llm_mlp.hpp" + +#include "transformations/itt.hpp" +namespace ov { + +template <> +EnumNames& EnumNames::get() { + static auto enum_names = EnumNames( + "op::intel_cpu::LLMMLPNode::ACT_FN", + {{"GELU", ov::intel_cpu::LLMMLPNode::ACT_FN::GELU}, {"SILU", ov::intel_cpu::LLMMLPNode::ACT_FN::SILU}}); + return enum_names; +} + +std::ostream& operator<<(std::ostream& os, const ov::intel_cpu::LLMMLPNode::ACT_FN& type) { + return os << as_string(type); +} + +namespace intel_cpu { + +bool LLMMLPNode::visit_attributes(ov::AttributeVisitor& visitor) { + INTERNAL_OP_SCOPE(LLMMLPNode_visit_attributes); + visitor.start_structure("config"); + visitor.on_attribute("act", m_config.act); + visitor.finish_structure(); + return true; +} + +void LLMMLPNode::validate_and_infer_types() { + INTERNAL_OP_SCOPE(LLMMLPNode_validate_and_infer_types); + const auto input_size = get_input_size(); + NODE_VALIDATION_CHECK(this, input_size == 4); + + const auto& ishape = get_input_partial_shape(0); + const auto& itype = get_input_element_type(0); + + const auto& w_down_shape = get_input_partial_shape(3); + + NODE_VALIDATION_CHECK(this, ishape.rank().is_static() && ishape.rank() == 3, "feature shape rank must be 3"); + const auto batch = ishape[0]; + const auto length = ishape[1]; + const auto feature = ishape[2]; + NODE_VALIDATION_CHECK(this, feature.is_static()); + + auto oshape = ishape; + oshape[oshape.size() - 1] = w_down_shape[0]; + set_output_type(0, itype, oshape); +} + +std::shared_ptr LLMMLPNode::clone_with_new_inputs(const ov::OutputVector& new_args) const { + INTERNAL_OP_SCOPE(LLMMLPNode_with_new_inputs); + check_new_args_count(this, new_args); + return std::make_shared(new_args, m_config); +} +} // namespace intel_cpu +} // namespace ov diff --git a/src/plugins/intel_cpu/src/transformations/cpu_opset/x64/op/llm_mlp.hpp b/src/plugins/intel_cpu/src/transformations/cpu_opset/x64/op/llm_mlp.hpp new file mode 100644 index 00000000000..f4f3c72b7e4 --- /dev/null +++ b/src/plugins/intel_cpu/src/transformations/cpu_opset/x64/op/llm_mlp.hpp @@ -0,0 +1,62 @@ +// Copyright (C) 2018-2024 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#pragma once + +#include "openvino/core/node.hpp" +#include "openvino/op/op.hpp" + +namespace ov { +namespace intel_cpu { + +class LLMMLPNode : public ov::op::Op { +public: + OPENVINO_OP("LLMMLP", "cpu_plugin_opset"); + + LLMMLPNode() = default; + + enum class ACT_FN { SILU = 0, GELU = 1}; + + struct Config { + ACT_FN act; + }; + + // args: + // 0: input + // 1: gate_proj + // 2: up_proj + // 3: down_proj + LLMMLPNode(const OutputVector& args, const Config& cfg) : Op(args), m_config(cfg) { + validate_and_infer_types(); + } + + bool visit_attributes(ov::AttributeVisitor& visitor) override; + + void validate_and_infer_types() override; + + std::shared_ptr clone_with_new_inputs(const ov::OutputVector& new_args) const override; + + const Config& get_config() const { + return m_config; + } + +private: + Config m_config; +}; + +} // namespace intel_cpu + +template <> +class AttributeAdapter + : public EnumAttributeAdapterBase { +public: + AttributeAdapter(ov::intel_cpu::LLMMLPNode::ACT_FN& value) + : EnumAttributeAdapterBase(value) {} + + OPENVINO_RTTI("AttributeAdapter"); +}; + +std::ostream& operator<<(std::ostream& s, const ov::intel_cpu::LLMMLPNode::ACT_FN& type); + +} // namespace ov diff --git a/src/plugins/intel_cpu/src/transformations/cpu_opset/x64/op/qkv_proj.cpp b/src/plugins/intel_cpu/src/transformations/cpu_opset/x64/op/qkv_proj.cpp new file mode 100644 index 00000000000..11d122b3815 --- /dev/null +++ b/src/plugins/intel_cpu/src/transformations/cpu_opset/x64/op/qkv_proj.cpp @@ -0,0 +1,39 @@ +// Copyright (C) 2018-2024 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include "qkv_proj.hpp" + +#include "transformations/itt.hpp" +namespace ov { +namespace intel_cpu { + +void QKVProjectionNode::validate_and_infer_types() { + INTERNAL_OP_SCOPE(QKVProjection_validate_and_infer_types); + const auto input_size = get_input_size(); + NODE_VALIDATION_CHECK(this, input_size == 4); + + const auto& ishape = get_input_partial_shape(0); + const auto& itype = get_input_element_type(0); + NODE_VALIDATION_CHECK(this, ishape.rank().is_static() && ishape.rank() == 3, "feature shape rank must be 3"); + + set_output_size(3); + + auto oshape0 = ishape; + auto oshape1 = ishape; + auto oshape2 = ishape; + oshape0[oshape0.size()-1] = get_input_partial_shape(1)[0]; + oshape1[oshape1.size()-1] = get_input_partial_shape(2)[0]; + oshape2[oshape2.size()-1] = get_input_partial_shape(3)[0]; + set_output_type(0, itype, oshape0); + set_output_type(1, itype, oshape1); + set_output_type(2, itype, oshape2); +} + +std::shared_ptr QKVProjectionNode::clone_with_new_inputs(const ov::OutputVector& new_args) const { + INTERNAL_OP_SCOPE(QKVProjection_with_new_inputs); + check_new_args_count(this, new_args); + return std::make_shared(new_args); +} +} // namespace intel_cpu +} // namespace ov diff --git a/src/plugins/intel_cpu/src/transformations/cpu_opset/x64/op/qkv_proj.hpp b/src/plugins/intel_cpu/src/transformations/cpu_opset/x64/op/qkv_proj.hpp new file mode 100644 index 00000000000..98f226ccdb4 --- /dev/null +++ b/src/plugins/intel_cpu/src/transformations/cpu_opset/x64/op/qkv_proj.hpp @@ -0,0 +1,34 @@ +// Copyright (C) 2018-2024 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#pragma once + +#include "openvino/core/node.hpp" +#include "openvino/op/op.hpp" + +namespace ov { +namespace intel_cpu { + +class QKVProjectionNode : public ov::op::Op { +public: + OPENVINO_OP("QKVProjection", "cpu_plugin_opset"); + + QKVProjectionNode() = default; + + // args: + // 0: input + // 1: gate_proj + // 2: up_proj + // 3: down_proj + QKVProjectionNode(const OutputVector& args) : Op(args) { + validate_and_infer_types(); + } + + void validate_and_infer_types() override; + + std::shared_ptr clone_with_new_inputs(const ov::OutputVector& new_args) const override; +}; + +} // namespace intel_cpu +} // namespace ov diff --git a/src/plugins/intel_cpu/src/transformations/cpu_opset/x64/pass/mlp_fusion.cpp b/src/plugins/intel_cpu/src/transformations/cpu_opset/x64/pass/mlp_fusion.cpp new file mode 100644 index 00000000000..d1efb226ff1 --- /dev/null +++ b/src/plugins/intel_cpu/src/transformations/cpu_opset/x64/pass/mlp_fusion.cpp @@ -0,0 +1,132 @@ +// Copyright (C) 2018-2024 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include "mlp_fusion.hpp" + +#include +#include +#include + +#include "itt.hpp" +#include "openvino/core/rt_info.hpp" +#include "openvino/opsets/opset1.hpp" +#include "openvino/opsets/opset6.hpp" +#include "openvino/opsets/opset7.hpp" +#include "openvino/opsets/opset8.hpp" +#include "openvino/pass/pattern/matcher.hpp" +#include "openvino/pass/pattern/op/or.hpp" +#include "openvino/pass/pattern/op/wrap_type.hpp" +#include "ov_ops/type_relaxed.hpp" +#include "transformations/cpu_opset/x64/op/llm_mlp.hpp" +#include "transformations/utils/utils.hpp" +#include "transformations/utils/gen_pattern.hpp" + +using namespace ov::gen_pattern; + +ov::intel_cpu::MLPFusion::MLPFusion() { + MATCHER_SCOPE(MLPFusion); + + auto input = makePattern("[?,?,?]"); + + auto gate_proj_weight_compressed = makePattern({}); // [up_size, down_size] + auto gate_proj_weight = makePattern({gate_proj_weight_compressed}, {{"destination_type", "f32"}}); + auto up_proj_weight_compressed = makePattern({}); // [up_size, down_size] + auto up_proj_weight = makePattern({up_proj_weight_compressed}, {{"destination_type", "f32"}}); + auto down_proj_weight_compressed = makePattern({}); // [down_size, up_size] + auto down_proj_weight = makePattern({down_proj_weight_compressed}, {{"destination_type", "f32"}}); + auto mlp_gate_proj = makePattern({input, gate_proj_weight | gate_proj_weight_compressed}, + {{"transpose_a", false}, {"transpose_b", true}}); // [?,?,up_size] + auto mlp_silu_gate = makePattern({mlp_gate_proj}); + auto mlp_gelu_gate = makePattern({mlp_gate_proj}); + auto mlp_up_proj = makePattern({input, up_proj_weight | up_proj_weight_compressed}, + {{"transpose_a", false}, {"transpose_b", true}}); + auto mlp_gated_up = makePattern({mlp_silu_gate | mlp_gelu_gate, mlp_up_proj}, {{"auto_broadcast", "numpy"}}); + auto down_proj = makePattern({mlp_gated_up, down_proj_weight | down_proj_weight_compressed}, + {{"transpose_a", false}, {"transpose_b", true}}); // [?,?,down_size] + + auto result = down_proj; + + matcher_pass_callback callback = [=](ov::pass::pattern::Matcher& m) { + PatternValidator validator(m); + if (!validator) { + return false; + } + + const auto& pattern_map = m.get_pattern_value_map(); + auto root = m.get_match_root(); + + auto gate_proj_w = pattern_map.at(gate_proj_weight_compressed); + auto up_proj_w = pattern_map.at(up_proj_weight_compressed); + auto down_proj_w = pattern_map.at(down_proj_weight_compressed); + + auto gate_proj_w_pshape = gate_proj_w.get_partial_shape(); + auto up_proj_w_pshape = up_proj_w.get_partial_shape(); + auto down_proj_w_pshape = down_proj_w.get_partial_shape(); + + // make sure that: + // - shape of gate/up's weight is [down_size, up_size] + // - shape of down's weight is [up_size, down_size] + if (!gate_proj_w_pshape.is_static()) + return false; + if (!up_proj_w_pshape.is_static()) + return false; + if (!down_proj_w_pshape.is_static()) + return false; + + auto up_shape = up_proj_w_pshape.get_shape(); + auto down_shape = down_proj_w_pshape.get_shape(); + + if (gate_proj_w_pshape.get_shape() != up_shape) + return false; + if (up_shape.size() != 2) + return false; + if (down_shape.size() != 2) + return false; + + auto up_size = up_shape[0]; + auto down_size = up_shape[1]; + if (down_shape[0] != down_size) + return false; + if (down_shape[1] != up_size) + return false; + + LLMMLPNode::Config config; + OutputVector new_args; + std::shared_ptr gate_act; + if (pattern_map.count(mlp_silu_gate) > 0) { + config.act = LLMMLPNode::ACT_FN::SILU; + gate_act = mlp_silu_gate; + } else if (pattern_map.count(mlp_gelu_gate) > 0) { + config.act = LLMMLPNode::ACT_FN::GELU; + gate_act = mlp_gelu_gate; + } else { + return false; + } + + new_args.push_back(pattern_map.at(input)); + new_args.push_back(gate_proj_w); + new_args.push_back(up_proj_w); + new_args.push_back(down_proj_w); + + auto old_node = root; + auto new_node = std::make_shared(new_args, config); + new_node->set_friendly_name(old_node->get_friendly_name()); + ov::copy_runtime_info({pattern_map.at(mlp_gate_proj).get_node_shared_ptr(), + pattern_map.at(gate_act).get_node_shared_ptr(), + pattern_map.at(mlp_up_proj).get_node_shared_ptr(), + pattern_map.at(down_proj).get_node_shared_ptr()}, + new_node); + + // callback is for plugin implementation to check if it can be supported + if (!transformation_callback(new_node)) { + return false; + } + + ov::replace_node(old_node, new_node); + return true; + }; + + auto m = std::make_shared(result, matcher_name); + this->register_matcher(m, callback); +} diff --git a/src/plugins/intel_cpu/src/transformations/cpu_opset/x64/pass/mlp_fusion.hpp b/src/plugins/intel_cpu/src/transformations/cpu_opset/x64/pass/mlp_fusion.hpp new file mode 100644 index 00000000000..aecb9292807 --- /dev/null +++ b/src/plugins/intel_cpu/src/transformations/cpu_opset/x64/pass/mlp_fusion.hpp @@ -0,0 +1,19 @@ +// Copyright (C) 2018-2024 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#pragma once + +#include + +namespace ov { +namespace intel_cpu { + +class MLPFusion: public ov::pass::MatcherPass { +public: + OPENVINO_RTTI("MLPFusion", "0"); + MLPFusion(); +}; + +} // namespace intel_cpu +} // namespace ov \ No newline at end of file diff --git a/src/plugins/intel_cpu/src/transformations/cpu_opset/x64/pass/qkv_proj_fusion.cpp b/src/plugins/intel_cpu/src/transformations/cpu_opset/x64/pass/qkv_proj_fusion.cpp new file mode 100644 index 00000000000..4b9540b06cc --- /dev/null +++ b/src/plugins/intel_cpu/src/transformations/cpu_opset/x64/pass/qkv_proj_fusion.cpp @@ -0,0 +1,123 @@ +// Copyright (C) 2018-2024 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include "qkv_proj_fusion.hpp" + +#include +#include +#include + +#include "itt.hpp" +#include "openvino/core/rt_info.hpp" +#include "openvino/opsets/opset1.hpp" +#include "openvino/opsets/opset6.hpp" +#include "openvino/opsets/opset8.hpp" +#include "openvino/pass/pattern/matcher.hpp" +#include "openvino/pass/pattern/op/or.hpp" +#include "openvino/pass/pattern/op/wrap_type.hpp" +#include "ov_ops/type_relaxed.hpp" +#include "transformations/cpu_opset/x64/op/qkv_proj.hpp" +#include "transformations/utils/utils.hpp" +#include "transformations/utils/gen_pattern.hpp" + +using namespace ov::gen_pattern; + +ov::intel_cpu::QKVProjFusion::QKVProjFusion() { + MATCHER_SCOPE(QKVProjFusion); + + auto input = makePattern("[?,?,?]"); + + auto q_proj_weight = makePattern({}); + auto q_proj_weight_cvt = + makePattern({q_proj_weight}, {{"destination_type", "f32"}}); // [4096,4096] + auto q_proj = makePattern({input, q_proj_weight_cvt | q_proj_weight}, + {{"transpose_a", false}, {"transpose_b", true}}); // [?,?,4096] + auto result = q_proj; + + matcher_pass_callback callback = [=](ov::pass::pattern::Matcher& m) { + PatternValidator validator(m); + if (!validator) { + return false; + } + + const auto& pattern_map = m.get_pattern_value_map(); + auto root = m.get_match_root(); + + auto src = pattern_map.at(input); + + auto&& children = src.get_target_inputs(); + + if (children.size() < 3) { + return false; + } + + OutputVector args = {src}; + OutputVector outputs; + size_t hidden_size = 0; + std::vector proj_size; + for (auto& child : children) { + auto mm = dynamic_cast(child.get_node()); + if (!mm) { + // maybe a ShapeOf + continue; + } + if (mm->get_transpose_a() != false || mm->get_transpose_b() != true) { + return false; + } + auto constw = ov::as_type_ptr(mm->input_value(1).get_node_shared_ptr()); + if (!constw) { + auto cvt = ov::as_type_ptr(mm->input_value(1).get_node_shared_ptr()); + if (!cvt) { + return false; + } + constw = ov::as_type_ptr(cvt->input_value(0).get_node_shared_ptr()); + } + if (!constw) { + return false; + } + + // input feature size should be the same + const auto& wshape = constw->get_shape(); + if (hidden_size == 0) { + hidden_size = wshape[1]; + } else if (hidden_size != wshape[1]) { + return false; + } + + proj_size.push_back(wshape[0]); + args.push_back(constw); + outputs.push_back(mm->get_default_output()); + } + + // make sure just 3 projections are found + if (outputs.size() != 3) { + return false; + } + if (args.size() != 4) { + return false; + } + + auto old_node = root; + auto new_node = std::make_shared(args); + new_node->set_friendly_name(old_node->get_friendly_name()); + ov::copy_runtime_info({old_node}, new_node); + + // callback is for plugin implementation to check if it can be supported + if (!transformation_callback(new_node)) { + return false; + } + + for (size_t i = 0; i < outputs.size(); i++) { + auto target = outputs[i].get_node_shared_ptr(); + outputs[i].replace(new_node->output(i)); + new_node->add_node_control_dependents(target); + new_node->add_node_control_dependencies(target); + target->clear_control_dependents(); + } + return true; + }; + + auto m = std::make_shared(result, matcher_name); + this->register_matcher(m, callback); +} diff --git a/src/plugins/intel_cpu/src/transformations/cpu_opset/x64/pass/qkv_proj_fusion.hpp b/src/plugins/intel_cpu/src/transformations/cpu_opset/x64/pass/qkv_proj_fusion.hpp new file mode 100644 index 00000000000..7127fdabf21 --- /dev/null +++ b/src/plugins/intel_cpu/src/transformations/cpu_opset/x64/pass/qkv_proj_fusion.hpp @@ -0,0 +1,19 @@ +// Copyright (C) 2018-2024 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#pragma once + +#include + +namespace ov { +namespace intel_cpu { + +class QKVProjFusion: public ov::pass::MatcherPass { +public: + OPENVINO_RTTI("QKVProjFusion", "0"); + QKVProjFusion(); +}; + +} // namespace intel_cpu +} // namespace ov \ No newline at end of file diff --git a/src/plugins/intel_cpu/src/transformations/transformation_pipeline.cpp b/src/plugins/intel_cpu/src/transformations/transformation_pipeline.cpp index 006935a85e8..7963a2bcb98 100644 --- a/src/plugins/intel_cpu/src/transformations/transformation_pipeline.cpp +++ b/src/plugins/intel_cpu/src/transformations/transformation_pipeline.cpp @@ -13,6 +13,7 @@ #include "openvino/opsets/opset5.hpp" #include "openvino/opsets/opset6.hpp" #include "openvino/opsets/opset10.hpp" +#include "openvino/op/paged_attention.hpp" #include #include #include @@ -115,6 +116,8 @@ #include "transformations/snippets/x64/pass/snippets_mark_skipped.hpp" #endif #include "transformations/cpu_opset/x64/pass/convert_to_interaction.hpp" +#include "transformations/cpu_opset/x64/pass/mlp_fusion.hpp" +#include "transformations/cpu_opset/x64/pass/qkv_proj_fusion.hpp" #include "transformations/cpu_opset/arm/pass/convert_group_conv.hpp" #include "transformations/cpu_opset/arm/pass/convert_group_conv1d.hpp" #include "transformations/cpu_opset/arm/pass/convert_reduce_multi_axis.hpp" @@ -143,6 +146,8 @@ #include "nodes/mha.h" #include "nodes/rnn.h" #include "nodes/scaled_attn.h" +#include "nodes/llm_mlp.h" +#include "nodes/qkv_proj.h" #include "dnnl.hpp" #if defined(OPENVINO_ARCH_ARM64) #include "cpu/aarch64/cpu_isa_traits.hpp" @@ -780,6 +785,37 @@ void Transformations::PostLpt() { CPU_REGISTER_PASS_X64(postLPTPassManager, ov::pass::RoPEFusion); CPU_REGISTER_PASS_X64(postLPTPassManager, CausalMaskPreprocessFusion); + // MLP & QKV fusion optimizations is focused on throughput, only enabled on AMX-bf16 & LLM serving use cases. + auto can_use_amx_bf16 = dnnl::impl::cpu::x64::mayiuse(dnnl::impl::cpu::x64::avx512_core_amx) && (inferencePrecision == element::bf16); + if (can_use_amx_bf16) { + auto has_paged_attention = op::util::has_op_with_type(model); + if (has_paged_attention) { + CPU_REGISTER_PASS_X64(postLPTPassManager, MLPFusion); + CPU_SET_CALLBACK_X64(postLPTPassManager, + [](const_node_ptr &node) -> bool { + std::string errorMsg; + return node::LLMMLP::isSupportedOperation(node, errorMsg); + }, + MLPFusion); + } + + // Limitations: at least 3 workers are required for QKV fusion + size_t concurrency = config.streamExecutorConfig.get_threads_per_stream(); + if (concurrency == 0) + concurrency = parallel_get_max_threads(); + if (concurrency >= 3) { + if (has_paged_attention) { + CPU_REGISTER_PASS_X64(postLPTPassManager, QKVProjFusion); + CPU_SET_CALLBACK_X64(postLPTPassManager, + [](const_node_ptr &node) -> bool { + std::string errorMsg; + return node::QKVProjection::isSupportedOperation(node, errorMsg); + }, + QKVProjFusion); + } + } + } + CPU_REGISTER_PASS_X64(postLPTPassManager, StatefulSDPAFusion); // Should be before Snippets pipeline because Ngram pattern contains eltwise nodes that can be tokenized by Snippets. diff --git a/src/plugins/intel_cpu/src/utils/plain_tensor.hpp b/src/plugins/intel_cpu/src/utils/plain_tensor.hpp index 9d2f787b3d0..9de9ba18c1e 100644 --- a/src/plugins/intel_cpu/src/utils/plain_tensor.hpp +++ b/src/plugins/intel_cpu/src/utils/plain_tensor.hpp @@ -119,6 +119,10 @@ struct PlainTensor { return m_strides[i]; } + size_t stride_bytes(int i) const { + return stride(i) * m_element_size; + } + template std::vector get_strides() const { std::vector strides(m_rank);