natas_hw
/
openvino
0
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[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 <cheng.luo@intel.com>
This commit is contained in:
Tingqian Li 2024-06-10 13:09:26 +08:00 committed by GitHub
parent a42da25019
commit f5463512c6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
23 changed files with 1956 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
src/plugins/intel_cpu/CMakeLists.txt
View File

@ -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})

4
src/plugins/intel_cpu/src/cpu_types.cpp
View File

@ -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

2
src/plugins/intel_cpu/src/cpu_types.h
View File

@ -124,6 +124,8 @@ enum class Type {
PagedAttention,
RoPE,
CausalMaskPreprocess,
LLMMLP,
QKVProjection,
};
enum class Algorithm {

4
src/plugins/intel_cpu/src/extension.cpp
View File

@ -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) \

385
src/plugins/intel_cpu/src/nodes/kernels/x64/mlp_kernel.cpp Normal file
View File

@ -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 <typename T>
void MKernel::repackB(ov::bfloat16* dst, T* src, int N_stride, int N, int K) {
if (N == 16 && K == 32 && std::is_same<T, ov::bfloat16>::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 <typename T>
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<ov::bfloat16>({static_cast<size_t>(N / 32), static_cast<size_t>(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<ov::bfloat16>(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<T>(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<T>(dst1, src1, N_stride, valid_n1, valid_k);
}
}
}
template void MKernel::prepareB<ov::bfloat16>(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<int>(repacked_B.size(0));
auto K = repacked_B.size(1) / 32;
auto* pB = repacked_B.ptr<uint8_t>();
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<jit_uni_eltwise_injector_f32<dnnl::impl::cpu::x64::avx512_core>>(
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

264
src/plugins/intel_cpu/src/nodes/kernels/x64/mlp_kernel.hpp Normal file
View File

@ -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 <typename T>
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 <typename T>
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<PlainTensor> weights; // ov::bfloat16 weights for current thread
std::shared_ptr<std::atomic_int> 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 <typename T>
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<float>({static_cast<size_t>(C_M), static_cast<size_t>(BN)});
auto pC = reinterpret_cast<uint8_t*>(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<uint8_t*>(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<uint8_t*>(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

28
src/plugins/intel_cpu/src/nodes/kernels/x64/mlp_utils.cpp Normal file
View File

@ -0,0 +1,28 @@
// Copyright (C) 2018-2024 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//
#include "mlp_utils.hpp"
#include <cstring>
#if defined(HAVE_AVX512F)
#include <immintrin.h>
#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<uint32_t*>(dst), reinterpret_cast<uint32_t*>(src), 16, stride/sizeof(uint32_t));
}
} // namespace XARCH
} // namespace Cpu
} // namespace Extensions
} // namespace ov

19
src/plugins/intel_cpu/src/nodes/kernels/x64/mlp_utils.hpp Normal file
View File

@ -0,0 +1,19 @@
// Copyright (C) 2018-2024 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//
#include <cstddef>
#include <vector>
#include <array>
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

339
src/plugins/intel_cpu/src/nodes/llm_mlp.cpp Normal file
View File

@ -0,0 +1,339 @@
// Copyright (C) 2018-2024 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//
#include "llm_mlp.h"
#include <string>
#include <vector>
#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<Work> 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 <typename T>
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<std::atomic_int>(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<void*>(p_weight));
}
void run(uint8_t* pA, int strideA, int M, ov::bfloat16* dstC, int strideC, std::vector<PlainTensor>& 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<float>(), peerC.ptr<float>(), workC.stride(0),
dstC + work.n0, strideC / sizeof(*dstC),
M, work.BN);
}
} else {
jit_reduce2bh_2.call(workC.ptr<float>(), 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<PlainTensor>& 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<float>(), 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<PlainTensor> 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<ov::bfloat16>({static_cast<size_t>(2 * N), static_cast<size_t>(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<ov::bfloat16>(), w_gate_up.stride_bytes(0), N * 2, K);
down.setup(w_down.ptr<ov::bfloat16>(), 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<size_t> scratch_offsets;
std::vector<size_t> 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<CpuBlockedMemoryDesc>(ov::element::u8, Shape{total_scratch_size});
m_scratchMem = m_scrachPad->createScratchPadMem(newMemDesc);
auto* scratch_base = m_scratchMem->getDataAs<uint8_t>();
m_actUp.resize<ov::bfloat16>({static_cast<size_t>(M), static_cast<size_t>(m_N)}, reinterpret_cast<ov::bfloat16*>(scratch_base));
for (size_t ithr = 0; ithr < m_tempC.size(); ithr++) {
m_tempC[ithr].resize<float>({1, scratch_C_sizes[ithr]}, reinterpret_cast<float*>(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<uint8_t>();
const auto& srcStrides = input->getDescWithType<BlockedMemoryDesc>()->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<ov::bfloat16>();
const auto& dstStrides = output->getDescWithType<BlockedMemoryDesc>()->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<ov::bfloat16>(), m_actUp.stride_bytes(0), m_config, m_tempC);
down.run(reinterpret_cast<uint8_t*>(m_actUp.ptr<ov::bfloat16>()), 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<ov::Node>& 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<const LLMMLPNode>(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<PortConfigurator> 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<PortConfigurator> 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<Impl>(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<const ov::Node>& op, std::string& errorMessage) noexcept {
try {
const auto node_mlp = std::dynamic_pointer_cast<const LLMMLPNode>(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

39
src/plugins/intel_cpu/src/nodes/llm_mlp.h Normal file
View File

@ -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<ov::Node>& 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<const ov::Node>& op, std::string& errorMessage) noexcept;
private:
struct Impl;
LLMMLPNode::Config m_mlp_config;
std::shared_ptr<Impl> m_pimpl;
};
} // namespace node
} // namespace intel_cpu
} // namespace ov

291
src/plugins/intel_cpu/src/nodes/qkv_proj.cpp Normal file
View File

@ -0,0 +1,291 @@
// Copyright (C) 2018-2024 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//
#include "qkv_proj.h"
#include <string>
#include <vector>
#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<int> allocate_workers(const std::vector<int>& grouped_works, int n_workers) {
auto n_groups = grouped_works.size();
// allocate 1 worker for each group
std::vector<int> 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<float>(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<Work> works;
std::vector<PlainTensor> 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<int>(w0.size(0));
auto proj_size1 = static_cast<int>(w1.size(0));
auto proj_size2 = static_cast<int>(w2.size(0));
auto n_group_workers = allocate_workers({proj_size0, proj_size1, proj_size2}, nthr);
create_works(w0.ptr<ov::bfloat16>(), 0, proj_size0, n_group_workers[0]);
create_works(w1.ptr<ov::bfloat16>(), 1, proj_size1, n_group_workers[1]);
create_works(w2.ptr<ov::bfloat16>(), 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<size_t> scratch_offsets;
std::vector<size_t> 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<CpuBlockedMemoryDesc>(ov::element::u8, Shape{total_scratch_size});
m_scratchMem = m_scrachPad->createScratchPadMem(newMemDesc);
auto* scratch_base = m_scratchMem->getDataAs<uint8_t>();
for (size_t ithr = 0; ithr < m_tempC.size(); ithr++) {
m_tempC[ithr].resize<float>({1, scratch_C_sizes[ithr]}, reinterpret_cast<float*>(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<uint8_t>();
int M = shape_size(ishape) / ishape[ishape.size() - 1];
auto* dst0 = m_node->getDstMemoryAtPort(0)->getDataAs<ov::bfloat16>();
auto* dst1 = m_node->getDstMemoryAtPort(1)->getDataAs<ov::bfloat16>();
auto* dst2 = m_node->getDstMemoryAtPort(2)->getDataAs<ov::bfloat16>();
const auto& srcStrides = input->getDescWithType<BlockedMemoryDesc>()->getStrides();
const auto& dstStrides0 = m_node->getDstMemoryAtPort(0)->getDescWithType<BlockedMemoryDesc>()->getStrides();
const auto& dstStrides1 = m_node->getDstMemoryAtPort(1)->getDescWithType<BlockedMemoryDesc>()->getStrides();
const auto& dstStrides2 = m_node->getDstMemoryAtPort(2)->getDescWithType<BlockedMemoryDesc>()->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<float>();
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<Impl>(this, context->getScratchPad());
}
}
void QKVProjection::execute(dnnl::stream strm) {
MAYBE_UNUSED(strm);
m_pimpl->execute();
}
QKVProjection::QKVProjection(const std::shared_ptr<ov::Node>& 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<PortConfigurator> 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<PortConfigurator> 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<const ov::Node>& op, std::string& errorMessage) noexcept {
try {
const auto node_qkv = std::dynamic_pointer_cast<const QKVProjectionNode>(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

41
src/plugins/intel_cpu/src/nodes/qkv_proj.h Normal file
View File

@ -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<ov::Node>& 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<const ov::Node>& op, std::string& errorMessage) noexcept;
private:
struct Impl;
std::shared_ptr<Impl> m_pimpl;
};
} // namespace node
} // namespace intel_cpu
} // namespace ov

4
src/plugins/intel_cpu/src/nodes_factory.cpp
View File

@ -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);

59
src/plugins/intel_cpu/src/transformations/cpu_opset/x64/op/llm_mlp.cpp Normal file
View File

@ -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<ov::intel_cpu::LLMMLPNode::ACT_FN>& EnumNames<ov::intel_cpu::LLMMLPNode::ACT_FN>::get() {
static auto enum_names = EnumNames<ov::intel_cpu::LLMMLPNode::ACT_FN>(
"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<Node> 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<LLMMLPNode>(new_args, m_config);
}
} // namespace intel_cpu
} // namespace ov

62
src/plugins/intel_cpu/src/transformations/cpu_opset/x64/op/llm_mlp.hpp Normal file
View File

@ -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<Node> 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<ov::intel_cpu::LLMMLPNode::ACT_FN>
: public EnumAttributeAdapterBase<ov::intel_cpu::LLMMLPNode::ACT_FN> {
public:
AttributeAdapter(ov::intel_cpu::LLMMLPNode::ACT_FN& value)
: EnumAttributeAdapterBase<ov::intel_cpu::LLMMLPNode::ACT_FN>(value) {}
OPENVINO_RTTI("AttributeAdapter<ov::intel_cpu::LLMMLPNode::ACT_FN>");
};
std::ostream& operator<<(std::ostream& s, const ov::intel_cpu::LLMMLPNode::ACT_FN& type);
} // namespace ov

39
src/plugins/intel_cpu/src/transformations/cpu_opset/x64/op/qkv_proj.cpp Normal file
View File

@ -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<Node> 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<QKVProjectionNode>(new_args);
}
} // namespace intel_cpu
} // namespace ov

34
src/plugins/intel_cpu/src/transformations/cpu_opset/x64/op/qkv_proj.hpp Normal file
View File

@ -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<Node> clone_with_new_inputs(const ov::OutputVector& new_args) const override;
};
} // namespace intel_cpu
} // namespace ov

132
src/plugins/intel_cpu/src/transformations/cpu_opset/x64/pass/mlp_fusion.cpp Normal file
View File

@ -0,0 +1,132 @@
// Copyright (C) 2018-2024 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//
#include "mlp_fusion.hpp"
#include <cstdint>
#include <iostream>
#include <limits>
#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<opset1::Constant>({}); // [up_size, down_size]
auto gate_proj_weight = makePattern<opset1::Convert>({gate_proj_weight_compressed}, {{"destination_type", "f32"}});
auto up_proj_weight_compressed = makePattern<opset1::Constant>({}); // [up_size, down_size]
auto up_proj_weight = makePattern<opset1::Convert>({up_proj_weight_compressed}, {{"destination_type", "f32"}});
auto down_proj_weight_compressed = makePattern<opset1::Constant>({}); // [down_size, up_size]
auto down_proj_weight = makePattern<opset1::Convert>({down_proj_weight_compressed}, {{"destination_type", "f32"}});
auto mlp_gate_proj = makePattern<opset1::MatMul>({input, gate_proj_weight | gate_proj_weight_compressed},
{{"transpose_a", false}, {"transpose_b", true}}); // [?,?,up_size]
auto mlp_silu_gate = makePattern<opset4::Swish>({mlp_gate_proj});
auto mlp_gelu_gate = makePattern<opset7::Gelu>({mlp_gate_proj});
auto mlp_up_proj = makePattern<opset1::MatMul>({input, up_proj_weight | up_proj_weight_compressed},
{{"transpose_a", false}, {"transpose_b", true}});
auto mlp_gated_up = makePattern<opset1::Multiply>({mlp_silu_gate | mlp_gelu_gate, mlp_up_proj}, {{"auto_broadcast", "numpy"}});
auto down_proj = makePattern<opset1::MatMul>({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<Node> 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<LLMMLPNode>(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<ov::pass::pattern::Matcher>(result, matcher_name);
this->register_matcher(m, callback);
}

19
src/plugins/intel_cpu/src/transformations/cpu_opset/x64/pass/mlp_fusion.hpp Normal file
View File

@ -0,0 +1,19 @@
// Copyright (C) 2018-2024 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//
#pragma once
#include <openvino/pass/graph_rewrite.hpp>
namespace ov {
namespace intel_cpu {
class MLPFusion: public ov::pass::MatcherPass {
public:
OPENVINO_RTTI("MLPFusion", "0");
MLPFusion();
};
} // namespace intel_cpu
} // namespace ov

123
src/plugins/intel_cpu/src/transformations/cpu_opset/x64/pass/qkv_proj_fusion.cpp Normal file
View File

@ -0,0 +1,123 @@
// Copyright (C) 2018-2024 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//
#include "qkv_proj_fusion.hpp"
#include <cstdint>
#include <iostream>
#include <limits>
#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<opset1::Constant>({});
auto q_proj_weight_cvt =
makePattern<opset1::Convert>({q_proj_weight}, {{"destination_type", "f32"}}); // [4096,4096]
auto q_proj = makePattern<opset1::MatMul>({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<int> proj_size;
for (auto& child : children) {
auto mm = dynamic_cast<opset1::MatMul*>(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<opset1::Constant>(mm->input_value(1).get_node_shared_ptr());
if (!constw) {
auto cvt = ov::as_type_ptr<opset1::Convert>(mm->input_value(1).get_node_shared_ptr());
if (!cvt) {
return false;
}
constw = ov::as_type_ptr<opset1::Constant>(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<QKVProjectionNode>(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<ov::pass::pattern::Matcher>(result, matcher_name);
this->register_matcher(m, callback);
}

19
src/plugins/intel_cpu/src/transformations/cpu_opset/x64/pass/qkv_proj_fusion.hpp Normal file
View File

@ -0,0 +1,19 @@
// Copyright (C) 2018-2024 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//
#pragma once
#include <openvino/pass/graph_rewrite.hpp>
namespace ov {
namespace intel_cpu {
class QKVProjFusion: public ov::pass::MatcherPass {
public:
OPENVINO_RTTI("QKVProjFusion", "0");
QKVProjFusion();
};
} // namespace intel_cpu
} // namespace ov

36
src/plugins/intel_cpu/src/transformations/transformation_pipeline.cpp
View File

@ -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 <ov_ops/augru_cell.hpp>
#include <ov_ops/augru_sequence.hpp>
#include <ov_ops/gather_compressed.hpp>
@ -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<ov::op::PagedAttentionExtension>(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.

4
src/plugins/intel_cpu/src/utils/plain_tensor.hpp
View File

@ -119,6 +119,10 @@ struct PlainTensor {
return m_strides[i];
}
size_t stride_bytes(int i) const {
return stride(i) * m_element_size;
}
template<typename T>
std::vector<T> get_strides() const {
std::vector<T> strides(m_rank);