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

[CPU] PagedAttention supports dynamic-split fuse (#24107) 

### Details:
 - *Merge first token and second token inference into one parallel loop*
- *~~Additional optimization: pre-transpose k-cache, pre-pack v-cache if
needed~~*
- *Additional optimization for first token: save q * k' upper triangle
matrix computation and (q * k') * v lower triangle matrix computation*
- *C++ pipeline can enable it:
https://github.com/ilya-lavrenov/openvino.genai/pull/9*
 - *TODO(in another PR):*
   - alibi support
   - performance tuning
   - testcase

### Tickets:
 - *[138673](https://jira.devtools.intel.com/browse/CVS-138673)*
This commit is contained in:
Luo Cheng 2024-05-15 16:46:14 +08:00 committed by GitHub
parent 6950460b34
commit 04a7ecfadf
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
24 changed files with 2817 additions and 1350 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
cmake/developer_package/cross_compile/cross_compiled_func.cmake
View File

@ -81,7 +81,8 @@ function(cross_compiled_file TARGET)
## that is arch ID
set(_arch ${_it})
if(_arch MATCHES ${_CURRENT_ARCH_FILTER})
list(APPEND _CUR_ARCH_SET ${_arch})
# make non/less-optimized version coming first
list(INSERT _CUR_ARCH_SET 0 ${_arch})
list(APPEND _FULL_ARCH_SET ${_arch})
endif()
else()

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

@ -209,6 +209,13 @@ cross_compiled_file(${TARGET_NAME}
NAME mha_single_token
NAMESPACE ov::Extensions::Cpu::XARCH
)
cross_compiled_file(${TARGET_NAME}
ARCH AVX512F AVX2 ANY
src/nodes/kernels/scaled_attn/executor_pa.cpp
API src/nodes/kernels/scaled_attn/executor_pa.hpp
NAME make_pa_executor
NAMESPACE ov::Extensions::Cpu::XARCH
)
cross_compiled_file(${TARGET_NAME}
ARCH AVX512F AVX2 ANY
src/nodes/kernels/scaled_attn/attn_memcpy.cpp

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

@ -237,7 +237,7 @@ static const TypeToNameMap& get_type_to_name_tbl() {
{"Ngram", Type::Ngram},
{"ScaledDotProductAttention", Type::ScaledDotProductAttention},
{"ScaledDotProductAttentionWithKVCache", Type::ScaledDotProductAttention},
{"PagedAttentionExtension", Type::ScaledDotProductAttention},
{"PagedAttentionExtension", Type::PagedAttention},
{"RoPE", Type::RoPE},
{"GatherCompressed", Type::Gather},
{"CausalMaskPreprocess", Type::CausalMaskPreprocess},
@ -358,6 +358,7 @@ std::string NameFromType(const Type type) {
CASE(Unique);
CASE(Ngram);
CASE(ScaledDotProductAttention);
CASE(PagedAttention);
CASE(RoPE);
CASE(CausalMaskPreprocess);
CASE(Unknown);

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

@ -118,6 +118,7 @@ enum class Type {
Unique,
Ngram,
ScaledDotProductAttention,
PagedAttention,
RoPE,
CausalMaskPreprocess,
};

3
src/plugins/intel_cpu/src/graph.cpp
View File

@ -1811,8 +1811,7 @@ void Graph::EnforceInferencePrecision() {
return true;
// kvcache of PagedAttention should be written directly
if (node->getType() == Type::ScaledDotProductAttention && node->getOriginalInputsNumber() == 13 &&
(inPort == 3 || inPort == 4))
if (node->getType() == Type::PagedAttention && (inPort == 3 || inPort == 4))
return true;
const auto &parent = node->getParentEdgeAt(inPort)->getParent();
/* Skip BF16 enforcement for nodes after Constant Inputs for maintaining precision for fusing.

33
src/plugins/intel_cpu/src/nodes/kernels/scaled_attn/attn_quant.cpp
View File

@ -17,6 +17,7 @@
#include "openvino/core/parallel.hpp"
#include "common.hpp"
#include "attn_quant.hpp"
#include "attn_quant_kernel.hpp"
namespace ov {
namespace Extensions {
@ -259,37 +260,7 @@ void attn_quant_u8(const float* src, uint8_t* dst, size_t n, float& scale, float
}
void attn_dequant_u8(const uint8_t* src, float* dst, size_t n, float scale, float zp) {
size_t i = 0;
// loadu_si128/epi64 does not support const qualifier
uint8_t* src_nc = const_cast<uint8_t*>(src);
#if defined(HAVE_AVX512F)
auto v_zp = _mm512_set1_ps(zp);
auto v_scale = _mm512_set1_ps(scale);
for (; i + vec_len_f32_avx512 <= n; i += vec_len_f32_avx512) {
auto v0_128 = _mm_loadu_si128(reinterpret_cast<__m128i*>(src_nc + i));
auto v0_512 = _mm512_cvtepu8_epi32(v0_128);
auto v0_value = _mm512_cvtepi32_ps(v0_512);
v0_value = _mm512_sub_ps(v0_value, v_zp);
auto v0_out = _mm512_mul_ps(v0_value, v_scale);
mm512_uni_storeu_ps(dst + i, v0_out);
}
#elif defined(HAVE_AVX2)
auto v_zp = _mm256_set1_ps(zp);
auto v_scale = _mm256_set1_ps(scale);
for (; i + vec_len_f32_avx2 <= n; i += vec_len_f32_avx2) {
auto v0_128 = _mm_loadl_epi64(reinterpret_cast<__m128i*>(src_nc + i));
auto v0_256 = _mm256_cvtepu8_epi32(v0_128);
auto v0_value = _mm256_cvtepi32_ps(v0_256);
v0_value = _mm256_sub_ps(v0_value, v_zp);
auto v0_out = _mm256_mul_ps(v0_value, v_scale);
mm256_uni_storeu_ps(dst + i, v0_out);
}
#endif
for (; i < n; ++i) {
float tmp = src_nc[i];
tmp = (tmp - zp) * scale;
dst[i] = tmp;
}
attn_dequant_u8_kernel(src, dst, n, scale, zp);
}
} // namespace XARCH

56
src/plugins/intel_cpu/src/nodes/kernels/scaled_attn/attn_quant_kernel.hpp Normal file
View File

@ -0,0 +1,56 @@
// Copyright (C) 2018-2024 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//
#pragma once
#include <array>
#include <cstddef>
#include <cstdint>
#include <vector>
#include "openvino/core/type/element_type.hpp"
#include "utils/plain_tensor.hpp"
namespace ov {
namespace Extensions {
namespace Cpu {
namespace XARCH {
template<typename TDST>
void attn_dequant_u8_kernel(const uint8_t* src, TDST* dst, size_t n, float scale, float zp) {
size_t i = 0;
// loadu_si128/epi64 does not support const qualifier
uint8_t* src_nc = const_cast<uint8_t*>(src);
#if defined(HAVE_AVX512F)
auto v_zp = _mm512_set1_ps(zp);
auto v_scale = _mm512_set1_ps(scale);
for (; i + vec_len_f32_avx512 <= n; i += vec_len_f32_avx512) {
auto v0_128 = _mm_loadu_si128(reinterpret_cast<__m128i*>(src_nc + i));
auto v0_512 = _mm512_cvtepu8_epi32(v0_128);
auto v0_value = _mm512_cvtepi32_ps(v0_512);
v0_value = _mm512_sub_ps(v0_value, v_zp);
auto v0_out = _mm512_mul_ps(v0_value, v_scale);
mm512_uni_storeu_ps(dst + i, v0_out);
}
#elif defined(HAVE_AVX2)
auto v_zp = _mm256_set1_ps(zp);
auto v_scale = _mm256_set1_ps(scale);
for (; i + vec_len_f32_avx2 <= n; i += vec_len_f32_avx2) {
auto v0_128 = _mm_loadl_epi64(reinterpret_cast<__m128i*>(src_nc + i));
auto v0_256 = _mm256_cvtepu8_epi32(v0_128);
auto v0_value = _mm256_cvtepi32_ps(v0_256);
v0_value = _mm256_sub_ps(v0_value, v_zp);
auto v0_out = _mm256_mul_ps(v0_value, v_scale);
mm256_uni_storeu_ps(dst + i, v0_out);
}
#endif
for (; i < n; ++i) {
float tmp = src_nc[i];
tmp = (tmp - zp) * scale;
dst[i] = tmp;
}
}
} // namespace XARCH
} // namespace Cpu
} // namespace Extensions
} // namespace ov

1619
src/plugins/intel_cpu/src/nodes/kernels/scaled_attn/executor_pa.cpp Normal file
View File

File diff suppressed because it is too large Load Diff

24
src/plugins/intel_cpu/src/nodes/kernels/scaled_attn/executor_pa.hpp Normal file
View File

@ -0,0 +1,24 @@
// Copyright (C) 2018-2024 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//
#pragma once
#include <array>
#include <cstddef>
#include <cstdint>
#include <vector>
#include <openvino/core/type/element_type.hpp>
#include "cpu_memory.h"
#include "executor_pa_common.hpp"
namespace ov {
namespace Extensions {
namespace Cpu {
namespace XARCH {
std::shared_ptr<PagedAttentionExecutor> make_pa_executor(ov::element::Type data_type, ov::element::Type kvcache_type);
} // namespace XARCH
} // namespace Cpu
} // namespace Extensions
} // namespace ov

113
src/plugins/intel_cpu/src/nodes/kernels/scaled_attn/executor_pa_common.cpp Normal file
View File

@ -0,0 +1,113 @@
// Copyright (C) 2018-2024 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//
#include <float.h>
#include <cmath>
#include <cstring>
#include <iostream>
#include <limits>
#include <type_traits>
#include "openvino/core/type/bfloat16.hpp"
#include "openvino/core/parallel.hpp"
#include "executor_pa_common.hpp"
#include "utils/plain_tensor.hpp"
namespace ov {
namespace Extensions {
namespace Cpu {
using namespace ov;
using namespace ov::intel_cpu;
using namespace dnnl::impl;
using namespace dnnl::impl::cpu::x64;
#ifdef OPENVINO_ARCH_X86_64
void TileConfig::reset(int palette, int _startRow, const std::vector<std::pair<int, int>>& _rows_columnsBytes) {
palette_id = palette;
startRow = _startRow;
unsigned long i;
for (i = 0; i < 14; i++) {
reserved[i] = 0;
}
for (i = 0; i < _rows_columnsBytes.size(); i++) {
rows[i] = _rows_columnsBytes[i].first;
cols[i] = _rows_columnsBytes[i].second;
}
for (; i < 16; i++) {
cols[i] = 0;
rows[i] = 0;
}
}
TileConfiger::TileConfiger() : jit_generator(jit_name()) {
create_kernel();
}
void TileConfiger::generate() {
Xbyak::Label release;
test(abi_param1, abi_param1);
jz(release);
ldtilecfg(ptr[abi_param1]);
ret();
L(release);
tilerelease();
ret();
}
JitMatMulVecAMX::JitMatMulVecAMX(int head_size, int block_size) : jit_generator(jit_name()), m_head_size(head_size), m_block_size(block_size) {
create_kernel();
m_tile_cfg.reset(1,
0,
{
{16, 4}, // C:0 M x 1 (4b)
{16, 64}, // A:1 M x 32/64 (64b)
{16, 4}, // B:2 32/64 x 1 (4b)
{16, 4}, // B:3
{16, 4}, // B:4
{16, 4}, // B:5
{16, 4}, // B:6
{16, 4}, // B:7
});
}
void JitMatMulVecAMX::generate() {
mov(reg_stride_A, m_head_size * 2);
mov(reg_stride_BC, 4);
const int kStep = 32;
if ((m_head_size % 32) != 0)
throw std::runtime_error("head size is not multiple of 32");
if ((m_block_size % 16) != 0)
throw std::runtime_error("block size is not multiple of 16");
auto num_B_tiles = m_head_size / kStep;
if (num_B_tiles > 6)
throw std::runtime_error("number of B tiles is bigger than 6");
/*
B(query) head_size x 1
A(key) matrix : block_size x head_size C(dst) block_size x 1
*/
// load query into B tiles
for (int i = 0; i < num_B_tiles; i++) {
tileloadd(Xbyak::Tmm(tmmB0.getIdx() + i), ptr[reg_q_addr + reg_stride_BC + i * 64]);
}
for (int m = 0; m < m_block_size; m += 16) {
tilezero(tmmC);
for (int i = 0; i < num_B_tiles; i++) {
tileloadd(tmmA, ptr[reg_k_addr + reg_stride_A + i * 64]);
tdpbf16ps(tmmC, tmmA, Xbyak::Tmm(tmmB0.getIdx() + i));
}
tilestored(ptr[reg_dst_addr + reg_stride_BC + m * sizeof(float)], tmmC);
add(reg_k_addr, m_head_size * 2 * 16);
}
ret();
}
#endif
} // namespace Cpu
} // namespace Extensions
} // namespace ov

107
src/plugins/intel_cpu/src/nodes/kernels/scaled_attn/executor_pa_common.hpp Normal file
View File

@ -0,0 +1,107 @@
// Copyright (C) 2018-2024 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//
#pragma once
#include <array>
#include <cstddef>
#include <cstdint>
#include <vector>
#include <openvino/core/type/element_type.hpp>
#include "cpu_memory.h"
#include "cpu/x64/cpu_isa_traits.hpp"
#include "cpu/x64/jit_generator.hpp"
namespace ov {
namespace Extensions {
namespace Cpu {
// this file will contain features that do not require multiple instantiation
struct PagedAttentionExecutor {
// PagedAttention input index
static const size_t ID_Q = 0;
static const size_t ID_K = 1;
static const size_t ID_V = 2;
static const size_t ID_KCACHE = 3;
static const size_t ID_VCACHE = 4;
static const size_t ID_IS_PROMPT = 5;
static const size_t ID_SLOT_MAPPING = 6;
static const size_t ID_MAX_CONTEXT_LEN = 7;
static const size_t ID_CONTEXT_LENS = 8;
static const size_t ID_BLOCK_TABLES = 9;
static const size_t ID_SCALE = 10;
static const size_t ID_ALIBI_SLOPES = 11;
static const size_t ID_SLIDING_WINDOW = 12;
static const size_t ID_SUBSEQUENCE_LENS = 13;
virtual void execute(const std::vector<ov::intel_cpu::MemoryPtr>& inputs, const ov::intel_cpu::MemoryPtr output) = 0;
};
#ifdef OPENVINO_ARCH_X86_64
// w = query * Key
//
// query: [1, S]
// Key : [block_size, S]
// w : [1, block_size]
//
// S is known at compile time
struct TileConfig {
uint8_t palette_id;
uint8_t startRow;
uint8_t reserved[14];
uint16_t cols[16];
uint8_t rows[16];
void reset(int palette, int _startRow, const std::vector<std::pair<int, int>>& _rows_columnsBytes);
};
class TileConfiger : public dnnl::impl::cpu::x64::jit_generator {
public:
DECLARE_CPU_JIT_AUX_FUNCTIONS(TileConfiger)
TileConfiger();
void generate() override;
};
class JitMatMulVecAMX : public dnnl::impl::cpu::x64::jit_generator {
void operator=(const JitMatMulVecAMX&);
public:
DECLARE_CPU_JIT_AUX_FUNCTIONS(JitMatMulVecAMX)
int m_head_size;
int m_block_size;
TileConfiger m_tile_configer;
TileConfig m_tile_cfg;
JitMatMulVecAMX(int head_size, int block_size);
void tile_config() {
m_tile_configer(&m_tile_cfg);
}
void tile_release() {
m_tile_configer(nullptr);
}
// to save push/pop: do not use `abi_save_gpr_regs`
static constexpr auto abi_param_regs = dnnl::impl::cpu::x64::abi_param_regs;
Xbyak::Reg64 reg_q_addr = abi_param1;
Xbyak::Reg64 reg_k_addr = abi_param2;
Xbyak::Reg64 reg_dst_addr = abi_param3;
Xbyak::Reg64 reg_stride_A = rax;
Xbyak::Reg64 reg_stride_BC = r9;
Xbyak::Tmm tmmC = tmm0;
Xbyak::Tmm tmmA = tmm1;
Xbyak::Tmm tmmB0 = tmm2;
Xbyak::Tmm tmmB1 = tmm3;
Xbyak::Tmm tmmB2 = tmm4;
Xbyak::Tmm tmmB3 = tmm5;
Xbyak::Tmm tmmB4 = tmm6;
Xbyak::Tmm tmmB5 = tmm7;
void generate() override;
};
#endif
} // namespace Cpu
} // namespace Extensions
} // namespace ov

892
src/plugins/intel_cpu/src/nodes/kernels/scaled_attn/mha_single_token.cpp
View File

@ -59,231 +59,6 @@ void cvt_copy(TA* dst, TB* src, size_t n) {
}
}
template<typename T>
static void attn_acc_value_block(float* out, float* weight, T* v, size_t S, size_t block_size) {
#if defined(HAVE_AVX512F)
size_t j = 0;
for (; j + 4 <= block_size; j += 4) {
auto attn_w_vec0 = _mm512_set1_ps(weight[0]);
auto attn_w_vec1 = _mm512_set1_ps(weight[1]);
auto attn_w_vec2 = _mm512_set1_ps(weight[2]);
auto attn_w_vec3 = _mm512_set1_ps(weight[3]);
size_t i = 0;
for (; i + vec_len_f32_avx512 <= S; i += vec_len_f32_avx512) {
auto v_out = mm512_uni_loadu_ps(out + i);
v_out = _mm512_fmadd_ps(attn_w_vec0, mm512_uni_loadu_ps(v + i), v_out);
v_out = _mm512_fmadd_ps(attn_w_vec1, mm512_uni_loadu_ps(v + i + S), v_out);
v_out = _mm512_fmadd_ps(attn_w_vec2, mm512_uni_loadu_ps(v + i + S * 2), v_out);
v_out = _mm512_fmadd_ps(attn_w_vec3, mm512_uni_loadu_ps(v + i + S * 3), v_out);
_mm512_storeu_ps(out + i, v_out);
}
for (; i < S; i++) {
out[i] += weight[0] * v[i];
out[i] += weight[1] * v[i + S];
out[i] += weight[2] * v[i + S * 2];
out[i] += weight[3] * v[i + S * 3];
}
v += 4 * S;
weight += 4;
}
if (j + 2 <= block_size) {
auto attn_w_vec0 = _mm512_set1_ps(weight[0]);
auto attn_w_vec1 = _mm512_set1_ps(weight[1]);
size_t i = 0;
for (; i + vec_len_f32_avx512 <= S; i += vec_len_f32_avx512) {
auto v_out = mm512_uni_loadu_ps(out + i);
v_out = _mm512_fmadd_ps(attn_w_vec0, mm512_uni_loadu_ps(v + i), v_out);
v_out = _mm512_fmadd_ps(attn_w_vec1, mm512_uni_loadu_ps(v + i + S), v_out);
_mm512_storeu_ps(out + i, v_out);
}
for (; i < S; i++) {
out[i] += weight[0] * v[i];
out[i] += weight[1] * v[i + S];
}
v += 2 * S;
weight += 2;
j += 2;
}
if (j < block_size) {
auto attn_w_vec0 = _mm512_set1_ps(weight[0]);
size_t i = 0;
for (; i + vec_len_f32_avx512 <= S; i += vec_len_f32_avx512) {
auto v_out = mm512_uni_loadu_ps(out + i);
v_out = _mm512_fmadd_ps(attn_w_vec0, mm512_uni_loadu_ps(v + i), v_out);
_mm512_storeu_ps(out + i, v_out);
}
for (; i < S; i++) {
out[i] += weight[0] * v[i];
}
}
return;
#elif defined(HAVE_AVX2)
size_t j = 0;
for (; j + 4 <= block_size; j += 4) {
auto attn_w_vec0 = _mm256_set1_ps(weight[0]);
auto attn_w_vec1 = _mm256_set1_ps(weight[1]);
auto attn_w_vec2 = _mm256_set1_ps(weight[2]);
auto attn_w_vec3 = _mm256_set1_ps(weight[3]);
size_t i = 0;
for (; i + vec_len_f32_avx2 <= S; i += vec_len_f32_avx2) {
auto v_out = mm256_uni_loadu_ps(out + i);
v_out = _mm256_fmadd_ps(attn_w_vec0, mm256_uni_loadu_ps(v + i), v_out);
v_out = _mm256_fmadd_ps(attn_w_vec1, mm256_uni_loadu_ps(v + i + S), v_out);
v_out = _mm256_fmadd_ps(attn_w_vec2, mm256_uni_loadu_ps(v + i + S * 2), v_out);
v_out = _mm256_fmadd_ps(attn_w_vec3, mm256_uni_loadu_ps(v + i + S * 3), v_out);
mm256_uni_storeu_ps(out + i, v_out);
}
for (; i < S; i++) {
out[i] += weight[0] * v[i];
out[i] += weight[1] * v[i + S];
out[i] += weight[2] * v[i + S * 2];
out[i] += weight[3] * v[i + S * 3];
}
v += 4 * S;
weight += 4;
}
if (j + 2 <= block_size) {
auto attn_w_vec0 = _mm256_set1_ps(weight[0]);
auto attn_w_vec1 = _mm256_set1_ps(weight[1]);
size_t i = 0;
for (; i + vec_len_f32_avx2 <= S; i += vec_len_f32_avx2) {
auto v_out = mm256_uni_loadu_ps(out + i);
v_out = _mm256_fmadd_ps(attn_w_vec0, mm256_uni_loadu_ps(v + i), v_out);
v_out = _mm256_fmadd_ps(attn_w_vec1, mm256_uni_loadu_ps(v + i + S), v_out);
mm256_uni_storeu_ps(out + i, v_out);
}
for (; i < S; i++) {
out[i] += weight[0] * v[i];
out[i] += weight[1] * v[i + S];
}
v += 2 * S;
weight += 2;
j += 2;
}
if (j < block_size) {
auto attn_w_vec0 = _mm256_set1_ps(weight[0]);
size_t i = 0;
for (; i + vec_len_f32_avx2 <= S; i += vec_len_f32_avx2) {
auto v_out = mm256_uni_loadu_ps(out + i);
v_out = _mm256_fmadd_ps(attn_w_vec0, mm256_uni_loadu_ps(v + i), v_out);
mm256_uni_storeu_ps(out + i, v_out);
}
for (; i < S; i++) {
out[i] += weight[0] * v[i];
}
}
return;
#endif
for (size_t j = 0; j < block_size; j++) {
for (size_t i = 0; i < S; i++) {
out[i] += weight[j] * v[i];
}
v += S;
}
}
static void attn_acc_value_block(float* out, float* weight, uint8_t* v, size_t S, size_t block_size) {
// The layout for per token per head:
// |scale(f32)|zeropoint(f32)|quantized feature(u8,idx_1)|quantized feature(u8,idx_2)|...|quantized feature(u8,idx_S)|
// The quantized feature will start from 8bytes=sizeof(float)+sizeof(float)
#if defined(HAVE_AVX512F)
size_t j = 0;
for (; j + 4 <= block_size; j += 4) {
auto v_f0 = reinterpret_cast<float*>(v);
auto v_f1 = reinterpret_cast<float*>(v + S + 8);
auto v_f2 = reinterpret_cast<float*>(v + 2 * (S + 8));
auto v_f3 = reinterpret_cast<float*>(v + 3 * (S + 8));
auto attn_w_vec0 = _mm512_set1_ps(weight[0] * v_f0[0]);
auto attn_w_vec1 = _mm512_set1_ps(weight[1] * v_f1[0]);
auto attn_w_vec2 = _mm512_set1_ps(weight[2] * v_f2[0]);
auto attn_w_vec3 = _mm512_set1_ps(weight[3] * v_f3[0]);
auto zp0 = _mm512_set1_ps(v_f0[1]);
auto zp1 = _mm512_set1_ps(v_f1[1]);
auto zp2 = _mm512_set1_ps(v_f2[1]);
auto zp3 = _mm512_set1_ps(v_f3[1]);
size_t i = 0;
v += 8;
for (; i + vec_len_f32_avx512 <= S; i += vec_len_f32_avx512) {
auto v_out = mm512_uni_loadu_ps(out + i);
auto v0 = _mm512_sub_ps(_mm512_cvtepi32_ps(_mm512_cvtepu8_epi32(_mm_loadu_si128(reinterpret_cast<__m128i*>(v + i)))), zp0);
auto v1 = _mm512_sub_ps(_mm512_cvtepi32_ps(_mm512_cvtepu8_epi32(_mm_loadu_si128(reinterpret_cast<__m128i*>(v + i + S + 8)))), zp1);
auto v2 = _mm512_sub_ps(_mm512_cvtepi32_ps(_mm512_cvtepu8_epi32(_mm_loadu_si128(reinterpret_cast<__m128i*>(v + i + 2 * (S + 8))))), zp2);
auto v3 = _mm512_sub_ps(_mm512_cvtepi32_ps(_mm512_cvtepu8_epi32(_mm_loadu_si128(reinterpret_cast<__m128i*>(v + i + 3 * (S + 8))))), zp3);
v_out = _mm512_fmadd_ps(attn_w_vec0, v0, v_out);
v_out = _mm512_fmadd_ps(attn_w_vec1, v1, v_out);
v_out = _mm512_fmadd_ps(attn_w_vec2, v2, v_out);
v_out = _mm512_fmadd_ps(attn_w_vec3, v3, v_out);
_mm512_storeu_ps(out + i, v_out);
}
for (; i < S; i++) {
out[i] += weight[0] * (v[i] - v_f0[1]) * v_f0[0];
out[i] += weight[1] * (v[i + S + 8] - v_f1[1]) * v_f1[0];
out[i] += weight[2] * (v[i + 2 * (S + 8)] - v_f2[1]) * v_f2[0];
out[i] += weight[3] * (v[i + 3 * (S + 8)] - v_f3[1]) * v_f3[0];
}
v += 4 * (S + 8) - 8;
weight += 4;
}
for (; j < block_size; j++) {
auto v_f0 = reinterpret_cast<float*>(v);
auto attn_w_vec0 = _mm512_set1_ps(weight[0] * v_f0[0]);
auto zp0 = _mm512_set1_ps(v_f0[1]);
size_t i = 0;
v += 8;
for (; i + vec_len_f32_avx512 <= S; i += vec_len_f32_avx512) {
auto v_out = mm512_uni_loadu_ps(out + i);
auto v0 = _mm512_sub_ps(_mm512_cvtepi32_ps(_mm512_cvtepu8_epi32(_mm_loadu_si128(reinterpret_cast<__m128i*>(v + i)))), zp0);
v_out = _mm512_fmadd_ps(attn_w_vec0, v0, v_out);
_mm512_storeu_ps(out + i, v_out);
}
for (; i < S; i++) {
out[i] += weight[0] * (v[i] - v_f0[1]) * v_f0[0];
}
v += S;
weight++;
}
return;
#elif defined(HAVE_AVX2)
size_t j = 0;
for (; j < block_size; j++) {
auto v_f0 = reinterpret_cast<float*>(v);
auto attn_w_vec0 = _mm256_set1_ps(weight[0] * v_f0[0]);
auto zp0 = _mm256_set1_ps(v_f0[1]);
size_t i = 0;
v += 8;
for (; i + vec_len_f32_avx2 <= S; i += vec_len_f32_avx2) {
auto v_out = mm256_uni_loadu_ps(out + i);
auto v0 = _mm256_sub_ps(_mm256_cvtepi32_ps(_mm256_cvtepu8_epi32(_mm_loadl_epi64(reinterpret_cast<__m128i*>(v + i)))), zp0);
v_out = _mm256_fmadd_ps(attn_w_vec0, v0, v_out);
mm256_uni_storeu_ps(out + i, v_out);
}
for (; i < S; i++) {
out[i] += weight[0] * (v[i] - v_f0[1]) * v_f0[0];
}
v += S;
weight++;
}
return;
#endif
for (size_t j = 0; j < block_size; j++) {
auto v0 = reinterpret_cast<float*>(v);
v += 8;
for (size_t i = 0; i < S; i++) {
out[i] += weight[j] * (v[i] - v0[1]) * v0[0];
}
v += S;
}
}
template<typename T>
static void attn_acc_value(float* out, float weight, T* v, size_t S, float* scale, float* zp) {
size_t i = 0;
@ -792,271 +567,6 @@ static float dot_product(TA* a, uint8_t* b, size_t n, float* scale, float* zp, f
#endif
}
template<typename TA, typename TB>
static void dot_product_block(TA* a, TB* b, float* c, size_t n, size_t block_size) {
#if defined(HAVE_AVX512F)
size_t j = 0;
for (; j + 4 <= block_size; j += 4) {
auto vsum0 = _mm512_setzero_ps();
auto vsum1 = _mm512_setzero_ps();
auto vsum2 = _mm512_setzero_ps();
auto vsum3 = _mm512_setzero_ps();
size_t i = 0;
for (; i + vec_len_f32_avx512 <= n; i += vec_len_f32_avx512) {
auto va = mm512_uni_loadu_ps(a + i);
vsum0 = _mm512_fmadd_ps(va, mm512_uni_loadu_ps(b + i), vsum0);
vsum1 = _mm512_fmadd_ps(va, mm512_uni_loadu_ps(b + i + n), vsum1);
vsum2 = _mm512_fmadd_ps(va, mm512_uni_loadu_ps(b + i + 2 * n), vsum2);
vsum3 = _mm512_fmadd_ps(va, mm512_uni_loadu_ps(b + i + 3 * n), vsum3);
}
float sum0 = _mm512_reduce_add_ps(vsum0);
float sum1 = _mm512_reduce_add_ps(vsum1);
float sum2 = _mm512_reduce_add_ps(vsum2);
float sum3 = _mm512_reduce_add_ps(vsum3);
for (; i < n; i++) {
sum0 += a[i] * b[i];
sum1 += a[i] * b[i + n];
sum2 += a[i] * b[i + 2 * n];
sum3 += a[i] * b[i + 3 * n];
}
c[0] = sum0;
c[1] = sum1;
c[2] = sum2;
c[3] = sum3;
c += 4;
b += 4 * n;
}
for (; j < block_size; j++) {
auto vsum = _mm512_setzero_ps();
size_t i = 0;
for (; i + vec_len_f32_avx512 <= n; i += vec_len_f32_avx512) {
auto va = mm512_uni_loadu_ps(a + i);
vsum = _mm512_fmadd_ps(va, mm512_uni_loadu_ps(b + i), vsum);
}
float sum = _mm512_reduce_add_ps(vsum);
for (; i < n; i++) {
sum += a[i] * b[i];
}
b += n;
*c++ = sum;
}
return;
#elif defined(HAVE_AVX2)
size_t j = 0;
for (; j + 4 <= block_size; j += 4) {
auto vsum0 = _mm256_set1_ps(0.0f);
auto vsum1 = _mm256_set1_ps(0.0f);
auto vsum2 = _mm256_set1_ps(0.0f);
auto vsum3 = _mm256_set1_ps(0.0f);
size_t i = 0;
for (; i + vec_len_f32_avx2 <= n; i += vec_len_f32_avx2) {
auto va = mm256_uni_loadu_ps(a + i);
vsum0 = _mm256_fmadd_ps(va, mm256_uni_loadu_ps(b + i), vsum0);
vsum1 = _mm256_fmadd_ps(va, mm256_uni_loadu_ps(b + i + n), vsum1);
vsum2 = _mm256_fmadd_ps(va, mm256_uni_loadu_ps(b + i + 2 * n), vsum2);
vsum3 = _mm256_fmadd_ps(va, mm256_uni_loadu_ps(b + i + 3 * n), vsum3);
}
hsum(vsum0);
hsum(vsum1);
hsum(vsum2);
hsum(vsum3);
float sum0 = _mm256_cvtss_f32(vsum0);
float sum1 = _mm256_cvtss_f32(vsum1);
float sum2 = _mm256_cvtss_f32(vsum2);
float sum3 = _mm256_cvtss_f32(vsum3);
for (; i < n; i++) {
sum0 += a[i] * b[i];
sum1 += a[i] * b[i + n];
sum2 += a[i] * b[i + 2 * n];
sum3 += a[i] * b[i + 3 * n];
}
c[0] = sum0;
c[1] = sum1;
c[2] = sum2;
c[3] = sum3;
c += 4;
b += 4 * n;
}
for (; j < block_size; j++) {
auto vsum = _mm256_set1_ps(0.0f);
size_t i = 0;
for (; i + vec_len_f32_avx2 <= n; i += vec_len_f32_avx2) {
auto va = mm256_uni_loadu_ps(a + i);
vsum = _mm256_fmadd_ps(va, mm256_uni_loadu_ps(b + i), vsum);
}
hsum(vsum);
float sum = _mm256_cvtss_f32(vsum);
for (; i < n; i++) {
sum += a[i] * b[i];
}
b += n;
*c++ = sum;
}
return;
#endif
for (size_t j = 0; j < block_size; j++) {
float sum = 0;
for (size_t i = 0; i < n; i++) {
sum += a[i] * b[i];
}
b += n;
*c++ = sum;
}
}
template<typename TA>
static void dot_product_block(TA* a, uint8_t* b, float* c, size_t n, size_t block_size) {
// The layout for per token per head:
// |scale(f32)|zeropoint(f32)|quantized feature(u8,idx_1)|quantized feature(u8,idx_2)|...|quantized feature(u8,idx_S)|
// The quantized feature will start from 8bytes=sizeof(float)+sizeof(float)
#if defined(HAVE_AVX512F)
size_t j = 0;
for (; j + 4 <= block_size; j += 4) {
auto vsum0 = _mm512_setzero_ps();
auto vsum1 = _mm512_setzero_ps();
auto vsum2 = _mm512_setzero_ps();
auto vsum3 = _mm512_setzero_ps();
auto b0 = reinterpret_cast<float*>(b);
auto b1 = reinterpret_cast<float*>(b + n + 8);
auto b2 = reinterpret_cast<float*>(b + (n + 8) * 2);
auto b3 = reinterpret_cast<float*>(b + (n + 8) * 3);
auto v_zp0 = _mm512_set1_ps(b0[1]);
auto v_zp1 = _mm512_set1_ps(b1[1]);
auto v_zp2 = _mm512_set1_ps(b2[1]);
auto v_zp3 = _mm512_set1_ps(b3[1]);
size_t i = 0;
b += 8;
for (; i + vec_len_f32_avx512 <= n; i += vec_len_f32_avx512) {
auto va = mm512_uni_loadu_ps(a + i);
auto vb0 = _mm512_sub_ps(_mm512_cvtepi32_ps(_mm512_cvtepu8_epi32(_mm_loadu_si128(reinterpret_cast<__m128i*>(b + i)))), v_zp0);
auto vb1 = _mm512_sub_ps(_mm512_cvtepi32_ps(_mm512_cvtepu8_epi32(_mm_loadu_si128(reinterpret_cast<__m128i*>(b + i + n + 8)))), v_zp1);
auto vb2 = _mm512_sub_ps(_mm512_cvtepi32_ps(_mm512_cvtepu8_epi32(_mm_loadu_si128(reinterpret_cast<__m128i*>(b + i + 2 * (n + 8))))), v_zp2);
auto vb3 = _mm512_sub_ps(_mm512_cvtepi32_ps(_mm512_cvtepu8_epi32(_mm_loadu_si128(reinterpret_cast<__m128i*>(b + i + 3 * (n + 8))))), v_zp3);
vsum0 = _mm512_fmadd_ps(va, vb0, vsum0);
vsum1 = _mm512_fmadd_ps(va, vb1, vsum1);
vsum2 = _mm512_fmadd_ps(va, vb2, vsum2);
vsum3 = _mm512_fmadd_ps(va, vb3, vsum3);
}
float sum0 = _mm512_reduce_add_ps(vsum0);
float sum1 = _mm512_reduce_add_ps(vsum1);
float sum2 = _mm512_reduce_add_ps(vsum2);
float sum3 = _mm512_reduce_add_ps(vsum3);
for (; i < n; i++) {
sum0 += a[i] * (b[i] - b0[1]);
sum1 += a[i] * (b[i + n + 8] - b1[1]);
sum2 += a[i] * (b[i + 2 * (n + 8)] - b2[1]);
sum3 += a[i] * (b[i + 3 * (n + 8)] - b3[1]);
}
c[0] = sum0 * b0[0];
c[1] = sum1 * b1[0];
c[2] = sum2 * b2[0];
c[3] = sum3 * b3[0];
c += 4;
b += 4 * (n + 8) - 8;
}
for (; j < block_size; j++) {
auto vsum = _mm512_setzero_ps();
auto b0 = reinterpret_cast<float*>(b);
auto v_zp = _mm512_set1_ps(b0[1]);
size_t i = 0;
b += 8;
for (; i + vec_len_f32_avx512 <= n; i += vec_len_f32_avx512) {
auto va = mm512_uni_loadu_ps(a + i);
auto vb = _mm512_sub_ps(_mm512_cvtepi32_ps(_mm512_cvtepu8_epi32(_mm_loadu_si128(reinterpret_cast<__m128i*>(b + i)))), v_zp);
vsum = _mm512_fmadd_ps(va, vb, vsum);
}
float sum = _mm512_reduce_add_ps(vsum);
for (; i < n; i++) {
sum += a[i] * (b[i] - b0[1]);
}
b += n;
*c++ = sum * b0[0];
}
return;
#elif defined(HAVE_AVX2)
size_t j = 0;
for (; j + 4 <= block_size; j += 4) {
auto vsum0 = _mm256_setzero_ps();
auto vsum1 = _mm256_setzero_ps();
auto vsum2 = _mm256_setzero_ps();
auto vsum3 = _mm256_setzero_ps();
auto b0 = reinterpret_cast<float*>(b);
auto b1 = reinterpret_cast<float*>(b + n + 8);
auto b2 = reinterpret_cast<float*>(b + (n + 8) * 2);
auto b3 = reinterpret_cast<float*>(b + (n + 8) * 3);
auto v_zp0 = _mm256_set1_ps(b0[1]);
auto v_zp1 = _mm256_set1_ps(b1[1]);
auto v_zp2 = _mm256_set1_ps(b2[1]);
auto v_zp3 = _mm256_set1_ps(b3[1]);
size_t i = 0;
b += 8;
for (; i + vec_len_f32_avx2 <= n; i += vec_len_f32_avx2) {
auto va = mm256_uni_loadu_ps(a + i);
auto vb0 = _mm256_sub_ps(_mm256_cvtepi32_ps(_mm256_cvtepu8_epi32(_mm_loadl_epi64(reinterpret_cast<__m128i*>(b + i)))), v_zp0);
auto vb1 = _mm256_sub_ps(_mm256_cvtepi32_ps(_mm256_cvtepu8_epi32(_mm_loadl_epi64(reinterpret_cast<__m128i*>(b + i + n + 8)))), v_zp1);
auto vb2 = _mm256_sub_ps(_mm256_cvtepi32_ps(_mm256_cvtepu8_epi32(_mm_loadl_epi64(reinterpret_cast<__m128i*>(b + i + 2 * (n + 8))))), v_zp2);
auto vb3 = _mm256_sub_ps(_mm256_cvtepi32_ps(_mm256_cvtepu8_epi32(_mm_loadl_epi64(reinterpret_cast<__m128i*>(b + i + 3 * (n + 8))))), v_zp3);
vsum0 = _mm256_fmadd_ps(va, vb0, vsum0);
vsum1 = _mm256_fmadd_ps(va, vb1, vsum1);
vsum2 = _mm256_fmadd_ps(va, vb2, vsum2);
vsum3 = _mm256_fmadd_ps(va, vb3, vsum3);
}
hsum(vsum0);
hsum(vsum1);
hsum(vsum2);
hsum(vsum3);
float sum0 = _mm256_cvtss_f32(vsum0);
float sum1 = _mm256_cvtss_f32(vsum1);
float sum2 = _mm256_cvtss_f32(vsum2);
float sum3 = _mm256_cvtss_f32(vsum3);
for (; i < n; i++) {
sum0 += a[i] * (b[i] - b0[1]);
sum1 += a[i] * (b[i + n + 8] - b1[1]);
sum2 += a[i] * (b[i + 2 * (n + 8)] - b2[1]);
sum3 += a[i] * (b[i + 3 * (n + 8)] - b3[1]);
}
c[0] = sum0 * b0[0];
c[1] = sum1 * b1[0];
c[2] = sum2 * b2[0];
c[3] = sum3 * b3[0];
c += 4;
b += 4 * (n + 8) - 8;
}
for (; j < block_size; j++) {
auto vsum = _mm256_setzero_ps();
auto b0 = reinterpret_cast<float*>(b);
auto v_zp = _mm256_set1_ps(b0[1]);
size_t i = 0;
b += 8;
for (; i + vec_len_f32_avx2 <= n; i += vec_len_f32_avx2) {
auto va = mm256_uni_loadu_ps(a + i);
auto vb = _mm256_sub_ps(_mm256_cvtepi32_ps(_mm256_cvtepu8_epi32(_mm_loadl_epi64(reinterpret_cast<__m128i*>(b + i)))), v_zp);
vsum = _mm256_fmadd_ps(va, vb, vsum);
}
hsum(vsum);
float sum = _mm256_cvtss_f32(vsum);
for (; i < n; i++) {
sum += a[i] * (b[i] - b0[1]);
}
b += n;
*c++ = sum * b0[0];
}
return;
#endif
for (size_t j = 0; j < block_size; j++) {
float sum = 0;
auto b0 = reinterpret_cast<float*>(b);
b += 8;
for (size_t i = 0; i < n; i++) {
sum += a[i] * (b[i] - b0[1]);
}
b += n;
*c++ = sum * b0[0];
}
}
template<typename T>
static void attn_reduce(T* dst, float* temp, size_t M, size_t S, size_t temp_stride) {
size_t i = 0;
@ -1103,8 +613,6 @@ static void mha_single_token_kernel(const ov::intel_cpu::PlainTensor& query,
const ov::intel_cpu::PlainTensor& alibi_mask,
const ov::intel_cpu::PlainTensor& attention_mask,
const ov::intel_cpu::PlainTensor& beams,
size_t max_context_len,
const ov::intel_cpu::PlainTensor& context_lens,
ov::intel_cpu::PlainTensor& output_emb,
ov::intel_cpu::PlainTensor& buf_attn_w,
ov::intel_cpu::PlainTensor& buf_attn_score,
@ -1122,27 +630,20 @@ static void mha_single_token_kernel(const ov::intel_cpu::PlainTensor& query,
auto S = query.size(3);
auto h_group_num = present_value.size(1);
size_t h_each_group_len = 1;
bool is_pagedattn = context_lens;
size_t block_size = present_value.size(2);
if (h_group_num != H) {
h_each_group_len = H / h_group_num;
}
if (d_scale == 0.0f)
d_scale = 1.0f / sqrt(S);
auto nthr = parallel_get_max_threads();
size_t kv_len;
if (is_pagedattn) {
kv_len = max_context_len;
} else {
kv_len = present_key.size(2);
}
auto kv_len = present_key.size(2);
#if defined(HAVE_AVX2) && !defined(HAVE_AVX512F)
// avx2 will pre-compute the zero point and try to save the sub instruction in the dot_product,
// but it seems not necessary for avx512. Possible reason may be that for avx2 the cost of dot_product
// is larger than the memory access time, but for avx512 is not and the cost of pre-compute is a pure increase.
bool pastkv_is_int8 = past_k_scale_zp;
if (pastkv_is_int8 && !is_pagedattn) {
if (pastkv_is_int8) {
// be sure no false sharing
head_sum.resize<float>({B, H, q_len, 16});
parallel_for3d(B, H, q_len, [&](size_t b, size_t h, size_t pq) {
@ -1151,271 +652,154 @@ static void mha_single_token_kernel(const ov::intel_cpu::PlainTensor& query,
}
#endif
// TODO: refactor to seperate files
if (is_pagedattn) {
// if present_key is true, it means q*k is already computed in the caller
if (present_key) {
if (B >= static_cast<size_t>(nthr)) {
parallel_for2d_dynamic(B, beams.m_dims[1], [&](size_t b, size_t pk_in_blocks) {
auto context_len = static_cast<size_t>(context_lens.ptr<int32_t>()[b]);
// kv_len must be valid
auto pk = pk_in_blocks * block_size;
if (pk < context_len) {
auto block_number = beams.ptr<int32_t>(b)[pk_in_blocks];
for (size_t h_group = 0; h_group < h_group_num; h_group++) {
for (size_t pq = 0; pq < q_len; pq++) {
for (size_t h = h_group * h_each_group_len; h < (h_group + 1) * h_each_group_len; h++) {
dot_product_block(query.ptr<T>(b, h, pq), present_key.ptr<T2>(block_number, h_group),
buf_attn_w.ptr<float>(b, h, pq) + pk, S, std::min(block_size, context_len - pk));
}
}
}
}
});
} else {
parallel_for3d_dynamic(B, beams.m_dims[1], h_group_num, [&](size_t b, size_t pk_in_blocks, size_t h_group) {
auto context_len = static_cast<size_t>(context_lens.ptr<int32_t>()[b]);
// kv_len must be valid
auto pk = pk_in_blocks * block_size;
if (pk < context_len) {
auto block_number = beams.ptr<int32_t>(b)[pk_in_blocks];
for (size_t pq = 0; pq < q_len; pq++) {
for (size_t h = h_group * h_each_group_len; h < (h_group + 1) * h_each_group_len; h++) {
dot_product_block(query.ptr<T>(b, h, pq), present_key.ptr<T2>(block_number, h_group),
buf_attn_w.ptr<float>(b, h, pq) + pk, S, std::min(block_size, context_len - pk));
}
}
}
});
}
}
parallel_nt_static(nthr, [&](const size_t ithr, const size_t nthr) {
size_t start{0}, end{0};
splitter(B * h_group_num * kv_len, nthr, ithr, start, end);
parallel_for3d_dynamic(B, H, q_len, [&](size_t b, size_t h, size_t pq) {
auto cur_kv_len = static_cast<size_t>(context_lens.ptr<int32_t>()[b]);
auto ncausal = cur_kv_len;
// apply attention mask & sofmax
float* alibi_ptr = alibi_mask ? &alibi_mask.at<float>({b, h, pq, 0}, true) : nullptr;
uint8_t* attn_mask_ptr = nullptr;
auto attn_mask_prec = attention_mask.get_precision();
if (attention_mask)
attn_mask_ptr = reinterpret_cast<uint8_t*>(&attention_mask.at<T>({b, h, pq, 0}, true));
uint8_t* cmask_ptr = causal_mask ? &causal_mask.at<uint8_t>({b, h, pq, 0}, true) : nullptr;
attn_softmax_kernel(buf_attn_w.ptr<float>(b, h, pq),
buf_attn_w.ptr<float>(b, h, pq),
d_scale,
alibi_ptr,
attn_mask_ptr,
cmask_ptr,
select_nfltmax_at_0,
ncausal,
cur_kv_len,
attn_mask_prec,
ov::element::f32);
});
// attn_w * V
// there are enough works for each thread
if (B >= static_cast<size_t>(nthr)) {
buf_attn_score.resize<float>({static_cast<size_t>(nthr), q_len, h_each_group_len, S});
parallel_for2d_dynamic(B, h_group_num, [&](size_t b, size_t h_group) {
auto ithr = parallel_get_thread_num();
auto context_len = static_cast<size_t>(context_lens.ptr<int32_t>()[b]);
memset(buf_attn_score.ptr<float>(ithr), 0, q_len * h_each_group_len * S * sizeof(float));
for (size_t pv = 0; pv < context_len; pv += block_size) {
size_t pv_in_blocks = pv / block_size;
auto block_number = beams.ptr<int32_t>(b)[pv_in_blocks];
auto* v = present_value.ptr<T2>(block_number, h_group);
for (size_t pq = 0; pq < q_len; pq++) {
for (size_t h = h_group * h_each_group_len, group_idx = 0; h < (h_group + 1) * h_each_group_len; h++, group_idx++) {
attn_acc_value_block(buf_attn_score.ptr<float>(ithr, pq, group_idx),
buf_attn_w.ptr<float>(b, h, pq) + pv,
v,
S,
std::min(block_size, context_len - pv));
}
}
}
// convert to dst
for (size_t pq = 0; pq < q_len; pq++)
for (size_t h = h_group * h_each_group_len, group_idx = 0; h < (h_group + 1) * h_each_group_len; h++, group_idx++)
cvt_copy(output_emb.ptr<T>(b, pq, h * S), buf_attn_score.ptr<float>(ithr, pq, group_idx), S);
});
return;
}
buf_attn_score.resize<float>({static_cast<size_t>(nthr), B, q_len, H, S});
// buf_attn_w {B, H, q_len, kv_len}
parallel_nt_static(nthr, [&](const size_t ithr, const size_t nthr) {
memset(buf_attn_score.ptr<float>(ithr, 0, 0, 0, 0), 0, buf_attn_score.stride(0) * sizeof(float));
});
auto kv_len_in_blocks = beams.m_dims[1];
parallel_for3d_dynamic(B, kv_len_in_blocks, h_group_num, [&](size_t b, size_t pv_in_blocks, size_t h_group) {
auto ithr = parallel_get_thread_num();
auto context_len = static_cast<size_t>(context_lens.ptr<int32_t>()[b]);
auto pv = pv_in_blocks * block_size;
// kv_len must be valid
if (pv < context_len) {
auto block_number = beams.ptr<int32_t>(b)[pv_in_blocks];
auto* v = present_value.ptr<T2>(block_number, h_group);
for (size_t pq = 0; pq < q_len; pq++) {
for (size_t h = h_group * h_each_group_len; h < (h_group + 1) * h_each_group_len; h++) {
attn_acc_value_block(buf_attn_score.ptr<float>(ithr, b, pq, h),
buf_attn_w.ptr<float>(b, h, pq) + pv,
v,
S,
std::min(block_size, context_len - pv));
}
}
}
});
} else {
parallel_nt_static(nthr, [&](const size_t ithr, const size_t nthr) {
size_t start{0}, end{0};
splitter(B * h_group_num * kv_len, nthr, ithr, start, end);
size_t b, h_group, pk;
if (start < end) {
parallel_it_init(start, pk, kv_len, b, B, h_group, h_group_num);
if (q_len == 1 && h_each_group_len == 1) {
if (B == 1) {
// the memory will be continuous when b==1
for (size_t iwork = start; iwork < end; ++iwork) {
auto p = past_k_scale_zp.ptr<float>(pk, 0, h_group);
auto p_k = present_key.ptr<T2>(0, h_group, pk);
prefetch_bytes(S, _MM_HINT_T0, 4096, p_k);
buf_attn_w.ptr<float>(0, h_group, 0)[pk] =
dot_product(query.ptr<T>(0, h_group), p_k,
S, p, p + 1, head_sum.ptr<float>(0, h_group));
parallel_it_step(pk, kv_len, b, B, h_group, h_group_num);
}
} else {
for (size_t iwork = start; iwork < end; ++iwork) {
auto b_kv = beams ? beams.ptr<int32_t>(b)[pk] : b;
auto p = past_k_scale_zp.ptr<float>(pk, b_kv, h_group);
auto p_k = present_key.ptr<T2>(b_kv, h_group, pk);
buf_attn_w.ptr<float>(b, h_group, 0)[pk] =
dot_product(query.ptr<T>(b, h_group), p_k,
S, p, p + 1, head_sum.ptr<float>(b, h_group));
parallel_it_step(pk, kv_len, b, B, h_group, h_group_num);
}
size_t b, h_group, pk;
if (start < end) {
parallel_it_init(start, pk, kv_len, b, B, h_group, h_group_num);
if (q_len == 1 && h_each_group_len == 1) {
if (B == 1) {
// the memory will be continuous when b==1
for (size_t iwork = start; iwork < end; ++iwork) {
auto p = past_k_scale_zp.ptr<float>(pk, 0, h_group);
auto p_k = present_key.ptr<T2>(0, h_group, pk);
prefetch_bytes(S, _MM_HINT_T0, 4096, p_k);
buf_attn_w.ptr<float>(0, h_group, 0)[pk] =
dot_product(query.ptr<T>(0, h_group), p_k,
S, p, p + 1, head_sum.ptr<float>(0, h_group));
parallel_it_step(pk, kv_len, b, B, h_group, h_group_num);
}
} else {
for (size_t iwork = start; iwork < end; ++iwork) {
auto b_kv = beams ? beams.ptr<int32_t>(b)[pk] : b;
for (size_t pq = 0; pq < q_len; pq++) {
auto p = past_k_scale_zp.ptr<float>(pk, b_kv, h_group);
for (size_t h = h_group * h_each_group_len; h < (h_group + 1) * h_each_group_len; h++) {
buf_attn_w.ptr<float>(b, h, pq)[pk] =
dot_product(query.ptr<T>(b, h, pq), present_key.ptr<T2>(b_kv, h_group, pk),
S, p, p + 1, head_sum.ptr<float>(b, h, pq));
}
}
auto p = past_k_scale_zp.ptr<float>(pk, b_kv, h_group);
auto p_k = present_key.ptr<T2>(b_kv, h_group, pk);
buf_attn_w.ptr<float>(b, h_group, 0)[pk] =
dot_product(query.ptr<T>(b, h_group), p_k,
S, p, p + 1, head_sum.ptr<float>(b, h_group));
parallel_it_step(pk, kv_len, b, B, h_group, h_group_num);
}
}
} else {
for (size_t iwork = start; iwork < end; ++iwork) {
auto b_kv = beams ? beams.ptr<int32_t>(b)[pk] : b;
for (size_t pq = 0; pq < q_len; pq++) {
auto p = past_k_scale_zp.ptr<float>(pk, b_kv, h_group);
for (size_t h = h_group * h_each_group_len; h < (h_group + 1) * h_each_group_len; h++) {
buf_attn_w.ptr<float>(b, h, pq)[pk] =
dot_product(query.ptr<T>(b, h, pq), present_key.ptr<T2>(b_kv, h_group, pk),
S, p, p + 1, head_sum.ptr<float>(b, h, pq));
}
}
parallel_it_step(pk, kv_len, b, B, h_group, h_group_num);
}
}
}
});
parallel_for3d(B, H, q_len, [&](size_t b, size_t h, size_t pq) {
auto cur_kv_len = kv_len;
auto ncausal = auto_causal ? (cur_kv_len - q_len + pq + 1) : cur_kv_len;
// apply attention mask & sofmax
float* alibi_ptr = alibi_mask ? &alibi_mask.at<float>({b, h, pq, 0}, true) : nullptr;
uint8_t* attn_mask_ptr = nullptr;
auto attn_mask_prec = attention_mask.get_precision();
if (attention_mask)
attn_mask_ptr = reinterpret_cast<uint8_t*>(&attention_mask.at<T>({b, h, pq, 0}, true));
uint8_t* cmask_ptr = causal_mask ? &causal_mask.at<uint8_t>({b, h, pq, 0}, true) : nullptr;
attn_softmax_kernel(buf_attn_w.ptr<float>(b, h, pq),
buf_attn_w.ptr<float>(b, h, pq),
d_scale,
alibi_ptr,
attn_mask_ptr,
cmask_ptr,
select_nfltmax_at_0,
ncausal,
cur_kv_len,
attn_mask_prec,
ov::element::f32);
});
// attn_w * V
// Fast Path if there are enough works for each thread
if (B >= static_cast<size_t>(nthr)) {
buf_attn_score.resize<float>({static_cast<size_t>(nthr), q_len, h_each_group_len, S});
parallel_for2d(B, h_group_num, [&](size_t b, size_t h_group) {
auto ithr = parallel_get_thread_num();
memset(buf_attn_score.ptr<float>(ithr), 0, q_len * h_each_group_len * S * sizeof(float));
for (size_t pv = 0; pv < kv_len; pv++) {
auto b_kv = beams ? beams.ptr<int32_t>(b)[pv] : b;
auto* v = present_value.ptr<T2>(b_kv, h_group, pv);
auto p = past_v_scale_zp.ptr<float>(pv, b_kv, h_group);
for (size_t pq = 0; pq < q_len; pq++) {
for (size_t h = h_group * h_each_group_len, group_idx = 0; h < (h_group + 1) * h_each_group_len; h++, group_idx++) {
attn_acc_value(buf_attn_score.ptr<float>(ithr, pq, group_idx),
buf_attn_w.ptr<float>(b, h, pq)[pv],
v,
S,
p + 0,
p + 1);
}
}
}
// convert to dst
for (size_t pq = 0; pq < q_len; pq++) {
for (size_t h = h_group * h_each_group_len, group_idx = 0; h < (h_group + 1) * h_each_group_len;
h++, group_idx++) {
auto* dst = has_out_transpose ? output_emb.ptr<T>(b, pq, h * S) : output_emb.ptr<T>(b, h, pq);
cvt_copy(dst, buf_attn_score.ptr<float>(ithr, pq, group_idx), S);
}
}
});
parallel_for3d(B, H, q_len, [&](size_t b, size_t h, size_t pq) {
auto cur_kv_len = kv_len;
auto ncausal = auto_causal ? (cur_kv_len - q_len + pq + 1) : cur_kv_len;
// apply attention mask & sofmax
float* alibi_ptr = alibi_mask ? &alibi_mask.at<float>({b, h, pq, 0}, true) : nullptr;
uint8_t* attn_mask_ptr = nullptr;
auto attn_mask_prec = attention_mask.get_precision();
if (attention_mask)
attn_mask_ptr = reinterpret_cast<uint8_t*>(&attention_mask.at<T>({b, h, pq, 0}, true));
uint8_t* cmask_ptr = causal_mask ? &causal_mask.at<uint8_t>({b, h, pq, 0}, true) : nullptr;
attn_softmax_kernel(buf_attn_w.ptr<float>(b, h, pq),
buf_attn_w.ptr<float>(b, h, pq),
d_scale,
alibi_ptr,
attn_mask_ptr,
cmask_ptr,
select_nfltmax_at_0,
ncausal,
cur_kv_len,
attn_mask_prec,
ov::element::f32);
});
// attn_w * V
// Fast Path if there are enough works for each thread
if (B >= static_cast<size_t>(nthr)) {
buf_attn_score.resize<float>({static_cast<size_t>(nthr), q_len, h_each_group_len, S});
parallel_for2d(B, h_group_num, [&](size_t b, size_t h_group) {
auto ithr = parallel_get_thread_num();
memset(buf_attn_score.ptr<float>(ithr), 0, q_len * h_each_group_len * S * sizeof(float));
for (size_t pv = 0; pv < kv_len; pv++) {
return;
}
buf_attn_score.resize<float>({static_cast<size_t>(nthr), B, q_len, H, S});
// buf_attn_w {B, H, q_len, kv_len}
parallel_nt_static(nthr, [&](const size_t ithr, const size_t nthr) {
size_t start{0}, end{0};
splitter(B * h_group_num * kv_len, nthr, ithr, start, end);
memset(buf_attn_score.ptr<float>(ithr, 0, 0, 0, 0), 0, buf_attn_score.stride(0) * sizeof(float));
size_t b, h_group, pv;
if (start < end) {
parallel_it_init(start, pv, kv_len, b, B, h_group, h_group_num);
if (q_len == 1 && h_each_group_len == 1) {
for (size_t iwork = start; iwork < end; ++iwork) {
auto b_kv = beams ? beams.ptr<int32_t>(b)[pv] : b;
auto* v = present_value.ptr<T2>(b_kv, h_group, pv);
auto p = past_v_scale_zp.ptr<float>(pv, b_kv, h_group);
attn_acc_value(buf_attn_score.ptr<float>(ithr, b, 0, h_group),
buf_attn_w.ptr<float>(b, h_group, 0, pv)[0],
v,
S,
p + 0,
p + 1);
parallel_it_step(pv, kv_len, b, B, h_group, h_group_num);
}
} else {
for (size_t iwork = start; iwork < end; ++iwork) {
auto b_kv = beams ? beams.ptr<int32_t>(b)[pv] : b;
auto* v = present_value.ptr<T2>(b_kv, h_group, pv);
auto p = past_v_scale_zp.ptr<float>(pv, b_kv, h_group);
for (size_t pq = 0; pq < q_len; pq++) {
for (size_t h = h_group * h_each_group_len, group_idx = 0; h < (h_group + 1) * h_each_group_len; h++, group_idx++) {
attn_acc_value(buf_attn_score.ptr<float>(ithr, pq, group_idx),
buf_attn_w.ptr<float>(b, h, pq)[pv],
v,
S,
p + 0,
p + 1);
for (size_t h = h_group * h_each_group_len; h < (h_group + 1) * h_each_group_len; h++) {
attn_acc_value(buf_attn_score.ptr<float>(ithr, b, pq, h),
buf_attn_w.ptr<float>(b, h, pq)[pv],
v,
S,
p + 0,
p + 1);
}
}
}
// convert to dst
for (size_t pq = 0; pq < q_len; pq++) {
for (size_t h = h_group * h_each_group_len, group_idx = 0; h < (h_group + 1) * h_each_group_len;
h++, group_idx++) {
auto* dst = has_out_transpose ? output_emb.ptr<T>(b, pq, h * S) : output_emb.ptr<T>(b, h, pq);
cvt_copy(dst, buf_attn_score.ptr<float>(ithr, pq, group_idx), S);
}
}
});
return;
}
buf_attn_score.resize<float>({static_cast<size_t>(nthr), B, q_len, H, S});
// buf_attn_w {B, H, q_len, kv_len}
parallel_nt_static(nthr, [&](const size_t ithr, const size_t nthr) {
size_t start{0}, end{0};
splitter(B * h_group_num * kv_len, nthr, ithr, start, end);
memset(buf_attn_score.ptr<float>(ithr, 0, 0, 0, 0), 0, buf_attn_score.stride(0) * sizeof(float));
size_t b, h_group, pv;
if (start < end) {
parallel_it_init(start, pv, kv_len, b, B, h_group, h_group_num);
if (q_len == 1 && h_each_group_len == 1) {
for (size_t iwork = start; iwork < end; ++iwork) {
auto b_kv = beams ? beams.ptr<int32_t>(b)[pv] : b;
auto* v = present_value.ptr<T2>(b_kv, h_group, pv);
auto p = past_v_scale_zp.ptr<float>(pv, b_kv, h_group);
attn_acc_value(buf_attn_score.ptr<float>(ithr, b, 0, h_group),
buf_attn_w.ptr<float>(b, h_group, 0, pv)[0],
v,
S,
p + 0,
p + 1);
parallel_it_step(pv, kv_len, b, B, h_group, h_group_num);
}
} else {
for (size_t iwork = start; iwork < end; ++iwork) {
auto b_kv = beams ? beams.ptr<int32_t>(b)[pv] : b;
auto* v = present_value.ptr<T2>(b_kv, h_group, pv);
auto p = past_v_scale_zp.ptr<float>(pv, b_kv, h_group);
for (size_t pq = 0; pq < q_len; pq++) {
for (size_t h = h_group * h_each_group_len; h < (h_group + 1) * h_each_group_len; h++) {
attn_acc_value(buf_attn_score.ptr<float>(ithr, b, pq, h),
buf_attn_w.ptr<float>(b, h, pq)[pv],
v,
S,
p + 0,
p + 1);
}
}
parallel_it_step(pv, kv_len, b, B, h_group, h_group_num);
}
parallel_it_step(pv, kv_len, b, B, h_group, h_group_num);
}
}
});
}
}
});
parallel_for3d(B, H, q_len, [&](size_t b, size_t h, size_t pq) {
auto* temp = buf_attn_score.ptr<float>(0, b, pq, h);
@ -1431,8 +815,6 @@ void mha_single_token(const ov::intel_cpu::PlainTensor& query,
const ov::intel_cpu::PlainTensor& alibi_mask,
const ov::intel_cpu::PlainTensor& attention_mask,
const ov::intel_cpu::PlainTensor& beams,
size_t max_context_len,
const ov::intel_cpu::PlainTensor& context_lens,
ov::intel_cpu::PlainTensor& output_emb,
ov::intel_cpu::PlainTensor& buf_attn_w,
ov::intel_cpu::PlainTensor& buf_attn_score,
@ -1450,8 +832,6 @@ void mha_single_token(const ov::intel_cpu::PlainTensor& query,
alibi_mask,
attention_mask,
beams,
max_context_len,
context_lens,
output_emb,
buf_attn_w,
buf_attn_score,
@ -1468,8 +848,6 @@ void mha_single_token(const ov::intel_cpu::PlainTensor& query,
alibi_mask,
attention_mask,
beams,
max_context_len,
context_lens,
output_emb,
buf_attn_w,
buf_attn_score,
@ -1488,8 +866,6 @@ void mha_single_token(const ov::intel_cpu::PlainTensor& query,
alibi_mask,
attention_mask,
beams,
max_context_len,
context_lens,
output_emb,
buf_attn_w,
buf_attn_score,
@ -1506,8 +882,6 @@ void mha_single_token(const ov::intel_cpu::PlainTensor& query,
alibi_mask,
attention_mask,
beams,
max_context_len,
context_lens,
output_emb,
buf_attn_w,
buf_attn_score,
@ -1524,8 +898,6 @@ void mha_single_token(const ov::intel_cpu::PlainTensor& query,
alibi_mask,
attention_mask,
beams,
max_context_len,
context_lens,
output_emb,
buf_attn_w,
buf_attn_score,

2
src/plugins/intel_cpu/src/nodes/kernels/scaled_attn/mha_single_token.hpp
View File

@ -21,8 +21,6 @@ void mha_single_token(const ov::intel_cpu::PlainTensor& query,
const ov::intel_cpu::PlainTensor& alibi_mask,
const ov::intel_cpu::PlainTensor& attention_mask,
const ov::intel_cpu::PlainTensor& beams,
size_t max_context_len,
const ov::intel_cpu::PlainTensor& context_lens,
ov::intel_cpu::PlainTensor& output_emb,
ov::intel_cpu::PlainTensor& buf_attn_w,
ov::intel_cpu::PlainTensor& buf_attn_score,

4
src/plugins/intel_cpu/src/nodes/kernels/scaled_attn/softmax_kernel.hpp
View File

@ -516,7 +516,7 @@ inline void attn_softmax_kernel(float* a,
ov::element::Type dst_precision) {
using func_fp32_type = void (*)(float*, float, const float*, const float*, const uint8_t*, bool, size_t, float&);
using func_bf16_type = void (*)(float*, float, const float*, const ov::bfloat16*, const uint8_t*, bool, size_t, float&);
static func_fp32_type funcs_fp32[] = {
static constexpr func_fp32_type funcs_fp32[] = {
scale_add2_reduce_max<false, false, false>,
scale_add2_reduce_max<false, false, true>,
scale_add2_reduce_max<false, true, false>,
@ -526,7 +526,7 @@ inline void attn_softmax_kernel(float* a,
scale_add2_reduce_max<true, true, false>,
scale_add2_reduce_max<true, true, true>
};
static func_bf16_type funcs_bf16[] = {
static constexpr func_bf16_type funcs_bf16[] = {
scale_add2_reduce_max<false, false, false>,
scale_add2_reduce_max<false, false, true>,
scale_add2_reduce_max<false, true, false>,

254
src/plugins/intel_cpu/src/nodes/kernels/scaled_attn/transpose_kernel.hpp Normal file
View File

@ -0,0 +1,254 @@
// Copyright (C) 2018-2024 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//
#pragma once
#include "common.hpp"
#include "openvino/core/type/element_type.hpp"
#include <array>
#include <cstddef>
#include <cstdint>
#include <vector>
namespace ov {
namespace Extensions {
namespace Cpu {
namespace XARCH {
#if defined(HAVE_AVX512F)
inline void transpose_m512i_16x16(__m512i& r0, __m512i& r1, __m512i& r2, __m512i& r3,
__m512i& r4, __m512i& r5, __m512i& r6, __m512i& r7,
__m512i& r8, __m512i& r9, __m512i& ra, __m512i& rb,
__m512i& rc, __m512i& rd, __m512i& re, __m512i& rf) {
__m512i t0, t1, t2, t3, t4, t5, t6, t7, t8, t9, ta, tb, tc, td, te, tf;
t0 = _mm512_unpacklo_epi32(r0, r1); // 0 16 1 17 4 20 5 21 8 24 9 25 12 28 13 29
t1 = _mm512_unpackhi_epi32(r0, r1); // 2 18 3 19 6 22 7 23 10 26 11 27 14 30 15 31
t2 = _mm512_unpacklo_epi32(r2, r3); // 32 48 33 49 ...
t3 = _mm512_unpackhi_epi32(r2, r3); // 34 50 35 51 ...
t4 = _mm512_unpacklo_epi32(r4, r5); // 64 80 65 81 ...
t5 = _mm512_unpackhi_epi32(r4, r5); // 66 82 67 83 ...
t6 = _mm512_unpacklo_epi32(r6, r7); // 96 112 97 113 ...
t7 = _mm512_unpackhi_epi32(r6, r7); // 98 114 99 115 ...
t8 = _mm512_unpacklo_epi32(r8, r9); // 128 ...
t9 = _mm512_unpackhi_epi32(r8, r9); // 130 ...
ta = _mm512_unpacklo_epi32(ra, rb); // 160 ...
tb = _mm512_unpackhi_epi32(ra, rb); // 162 ...
tc = _mm512_unpacklo_epi32(rc, rd); // 196 ...
td = _mm512_unpackhi_epi32(rc, rd); // 198 ...
te = _mm512_unpacklo_epi32(re, rf); // 228 ...
tf = _mm512_unpackhi_epi32(re, rf); // 230 ...
r0 = _mm512_unpacklo_epi64(t0, t2); // 0 16 32 48 ...
r1 = _mm512_unpackhi_epi64(t0, t2); // 1 17 33 49 ...
r2 = _mm512_unpacklo_epi64(t1, t3); // 2 18 34 49 ...
r3 = _mm512_unpackhi_epi64(t1, t3); // 3 19 35 51 ...
r4 = _mm512_unpacklo_epi64(t4, t6); // 64 80 96 112 ...
r5 = _mm512_unpackhi_epi64(t4, t6); // 65 81 97 114 ...
r6 = _mm512_unpacklo_epi64(t5, t7); // 66 82 98 113 ...
r7 = _mm512_unpackhi_epi64(t5, t7); // 67 83 99 115 ...
r8 = _mm512_unpacklo_epi64(t8, ta); // 128 144 160 176 ...
r9 = _mm512_unpackhi_epi64(t8, ta); // 129 145 161 178 ...
ra = _mm512_unpacklo_epi64(t9, tb); // 130 146 162 177 ...
rb = _mm512_unpackhi_epi64(t9, tb); // 131 147 163 179 ...
rc = _mm512_unpacklo_epi64(tc, te); // 192 208 228 240 ...
rd = _mm512_unpackhi_epi64(tc, te); // 193 209 229 241 ...
re = _mm512_unpacklo_epi64(td, tf); // 194 210 230 242 ...
rf = _mm512_unpackhi_epi64(td, tf); // 195 211 231 243 ...
t0 = _mm512_shuffle_i32x4(r0, r4, 0x88); // 0 16 32 48 8 24 40 56 64 80 96 112 ...
t1 = _mm512_shuffle_i32x4(r1, r5, 0x88); // 1 17 33 49 ...
t2 = _mm512_shuffle_i32x4(r2, r6, 0x88); // 2 18 34 50 ...
t3 = _mm512_shuffle_i32x4(r3, r7, 0x88); // 3 19 35 51 ...
t4 = _mm512_shuffle_i32x4(r0, r4, 0xdd); // 4 20 36 52 ...
t5 = _mm512_shuffle_i32x4(r1, r5, 0xdd); // 5 21 37 53 ...
t6 = _mm512_shuffle_i32x4(r2, r6, 0xdd); // 6 22 38 54 ...
t7 = _mm512_shuffle_i32x4(r3, r7, 0xdd); // 7 23 39 55 ...
t8 = _mm512_shuffle_i32x4(r8, rc, 0x88); // 128 144 160 176 ...
t9 = _mm512_shuffle_i32x4(r9, rd, 0x88); // 129 145 161 177 ...
ta = _mm512_shuffle_i32x4(ra, re, 0x88); // 130 146 162 178 ...
tb = _mm512_shuffle_i32x4(rb, rf, 0x88); // 131 147 163 179 ...
tc = _mm512_shuffle_i32x4(r8, rc, 0xdd); // 132 148 164 180 ...
td = _mm512_shuffle_i32x4(r9, rd, 0xdd); // 133 149 165 181 ...
te = _mm512_shuffle_i32x4(ra, re, 0xdd); // 134 150 166 182 ...
tf = _mm512_shuffle_i32x4(rb, rf, 0xdd); // 135 151 167 183 ...
r0 = _mm512_shuffle_i32x4(t0, t8, 0x88); // 0 16 32 48 64 80 96 112 ... 240
r1 = _mm512_shuffle_i32x4(t1, t9, 0x88); // 1 17 33 49 66 81 97 113 ... 241
r2 = _mm512_shuffle_i32x4(t2, ta, 0x88); // 2 18 34 50 67 82 98 114 ... 242
r3 = _mm512_shuffle_i32x4(t3, tb, 0x88); // 3 19 35 51 68 83 99 115 ... 243
r4 = _mm512_shuffle_i32x4(t4, tc, 0x88); // 4 ...
r5 = _mm512_shuffle_i32x4(t5, td, 0x88); // 5 ...
r6 = _mm512_shuffle_i32x4(t6, te, 0x88); // 6 ...
r7 = _mm512_shuffle_i32x4(t7, tf, 0x88); // 7 ...
r8 = _mm512_shuffle_i32x4(t0, t8, 0xdd); // 8 ...
r9 = _mm512_shuffle_i32x4(t1, t9, 0xdd); // 9 ...
ra = _mm512_shuffle_i32x4(t2, ta, 0xdd); // 10 ...
rb = _mm512_shuffle_i32x4(t3, tb, 0xdd); // 11 ...
rc = _mm512_shuffle_i32x4(t4, tc, 0xdd); // 12 ...
rd = _mm512_shuffle_i32x4(t5, td, 0xdd); // 13 ...
re = _mm512_shuffle_i32x4(t6, te, 0xdd); // 14 ...
rf = _mm512_shuffle_i32x4(t7, tf, 0xdd); // 15 31 47 63 79 96 111 127 ... 255
}
template<typename T>
inline void transpose_16x16_kernel(float* _dst, T* src, size_t dst_stride, size_t src_stride) {
auto* dst = reinterpret_cast<uint32_t*>(_dst);
__m512i r0, r1, r2, r3, r4, r5, r6, r7, r8, r9, ra, rb, rc, rd, re, rf;
r0 = _mm512_castps_si512(mm512_uni_loadu_ps(src));
r1 = _mm512_castps_si512(mm512_uni_loadu_ps(src + src_stride));
r2 = _mm512_castps_si512(mm512_uni_loadu_ps(src + 2 * src_stride));
r3 = _mm512_castps_si512(mm512_uni_loadu_ps(src + 3 * src_stride));
r4 = _mm512_castps_si512(mm512_uni_loadu_ps(src + 4 * src_stride));
r5 = _mm512_castps_si512(mm512_uni_loadu_ps(src + 5 * src_stride));
r6 = _mm512_castps_si512(mm512_uni_loadu_ps(src + 6 * src_stride));
r7 = _mm512_castps_si512(mm512_uni_loadu_ps(src + 7 * src_stride));
r8 = _mm512_castps_si512(mm512_uni_loadu_ps(src + 8 * src_stride));
r9 = _mm512_castps_si512(mm512_uni_loadu_ps(src + 9 * src_stride));
ra = _mm512_castps_si512(mm512_uni_loadu_ps(src + 10 * src_stride));
rb = _mm512_castps_si512(mm512_uni_loadu_ps(src + 11 * src_stride));
rc = _mm512_castps_si512(mm512_uni_loadu_ps(src + 12 * src_stride));
rd = _mm512_castps_si512(mm512_uni_loadu_ps(src + 13 * src_stride));
re = _mm512_castps_si512(mm512_uni_loadu_ps(src + 14 * src_stride));
rf = _mm512_castps_si512(mm512_uni_loadu_ps(src + 15 * src_stride));
transpose_m512i_16x16(r0, r1, r2, r3, r4, r5, r6, r7, r8, r9, ra, rb, rc, rd, re, rf);
_mm512_storeu_si512(dst, r0);
_mm512_storeu_si512(dst + dst_stride, r1);
_mm512_storeu_si512(dst + 2 * dst_stride, r2);
_mm512_storeu_si512(dst + 3 * dst_stride, r3);
_mm512_storeu_si512(dst + 4 * dst_stride, r4);
_mm512_storeu_si512(dst + 5 * dst_stride, r5);
_mm512_storeu_si512(dst + 6 * dst_stride, r6);
_mm512_storeu_si512(dst + 7 * dst_stride, r7);
_mm512_storeu_si512(dst + 8 * dst_stride, r8);
_mm512_storeu_si512(dst + 9 * dst_stride, r9);
_mm512_storeu_si512(dst + 10 * dst_stride, ra);
_mm512_storeu_si512(dst + 11 * dst_stride, rb);
_mm512_storeu_si512(dst + 12 * dst_stride, rc);
_mm512_storeu_si512(dst + 13 * dst_stride, rd);
_mm512_storeu_si512(dst + 14 * dst_stride, re);
_mm512_storeu_si512(dst + 15 * dst_stride, rf);
}
inline void transpose_16x16_kernel(uint32_t* dst, uint32_t* src, size_t dst_stride, size_t src_stride) {
__m512i r0, r1, r2, r3, r4, r5, r6, r7, r8, r9, ra, rb, rc, rd, re, rf;
r0 = _mm512_loadu_si512(src);
r1 = _mm512_loadu_si512(src + src_stride);
r2 = _mm512_loadu_si512(src + 2 * src_stride);
r3 = _mm512_loadu_si512(src + 3 * src_stride);
r4 = _mm512_loadu_si512(src + 4 * src_stride);
r5 = _mm512_loadu_si512(src + 5 * src_stride);
r6 = _mm512_loadu_si512(src + 6 * src_stride);
r7 = _mm512_loadu_si512(src + 7 * src_stride);
r8 = _mm512_loadu_si512(src + 8 * src_stride);
r9 = _mm512_loadu_si512(src + 9 * src_stride);
ra = _mm512_loadu_si512(src + 10 * src_stride);
rb = _mm512_loadu_si512(src + 11 * src_stride);
rc = _mm512_loadu_si512(src + 12 * src_stride);
rd = _mm512_loadu_si512(src + 13 * src_stride);
re = _mm512_loadu_si512(src + 14 * src_stride);
rf = _mm512_loadu_si512(src + 15 * src_stride);
transpose_m512i_16x16(r0, r1, r2, r3, r4, r5, r6, r7, r8, r9, ra, rb, rc, rd, re, rf);
_mm512_storeu_si512(dst, r0);
_mm512_storeu_si512(dst + dst_stride, r1);
_mm512_storeu_si512(dst + 2 * dst_stride, r2);
_mm512_storeu_si512(dst + 3 * dst_stride, r3);
_mm512_storeu_si512(dst + 4 * dst_stride, r4);
_mm512_storeu_si512(dst + 5 * dst_stride, r5);
_mm512_storeu_si512(dst + 6 * dst_stride, r6);
_mm512_storeu_si512(dst + 7 * dst_stride, r7);
_mm512_storeu_si512(dst + 8 * dst_stride, r8);
_mm512_storeu_si512(dst + 9 * dst_stride, r9);
_mm512_storeu_si512(dst + 10 * dst_stride, ra);
_mm512_storeu_si512(dst + 11 * dst_stride, rb);
_mm512_storeu_si512(dst + 12 * dst_stride, rc);
_mm512_storeu_si512(dst + 13 * dst_stride, rd);
_mm512_storeu_si512(dst + 14 * dst_stride, re);
_mm512_storeu_si512(dst + 15 * dst_stride, rf);
}
#elif defined(HAVE_AVX2)
// https://stackoverflow.com/questions/25622745/transpose-an-8x8-float-using-avx-avx2
inline void transpose_8x8(__m256& r0, __m256& r1, __m256& r2, __m256& r3, __m256& r4, __m256& r5, __m256& r6, __m256& r7) {
__m256 t0, t1, t2, t3, t4, t5, t6, t7;
__m256 tt0, tt1, tt2, tt3, tt4, tt5, tt6, tt7;
t0 = _mm256_unpacklo_ps(r0, r1);
t1 = _mm256_unpackhi_ps(r0, r1);
t2 = _mm256_unpacklo_ps(r2, r3);
t3 = _mm256_unpackhi_ps(r2, r3);
t4 = _mm256_unpacklo_ps(r4, r5);
t5 = _mm256_unpackhi_ps(r4, r5);
t6 = _mm256_unpacklo_ps(r6, r7);
t7 = _mm256_unpackhi_ps(r6, r7);
tt0 = _mm256_shuffle_ps(t0, t2, _MM_SHUFFLE(1, 0, 1, 0));
tt1 = _mm256_shuffle_ps(t0, t2, _MM_SHUFFLE(3, 2, 3, 2));
tt2 = _mm256_shuffle_ps(t1, t3, _MM_SHUFFLE(1, 0, 1, 0));
tt3 = _mm256_shuffle_ps(t1, t3, _MM_SHUFFLE(3, 2, 3, 2));
tt4 = _mm256_shuffle_ps(t4, t6, _MM_SHUFFLE(1, 0, 1, 0));
tt5 = _mm256_shuffle_ps(t4, t6, _MM_SHUFFLE(3, 2, 3, 2));
tt6 = _mm256_shuffle_ps(t5, t7, _MM_SHUFFLE(1, 0, 1, 0));
tt7 = _mm256_shuffle_ps(t5, t7, _MM_SHUFFLE(3, 2, 3, 2));
r0 = _mm256_permute2f128_ps(tt0, tt4, 0x20);
r1 = _mm256_permute2f128_ps(tt1, tt5, 0x20);
r2 = _mm256_permute2f128_ps(tt2, tt6, 0x20);
r3 = _mm256_permute2f128_ps(tt3, tt7, 0x20);
r4 = _mm256_permute2f128_ps(tt0, tt4, 0x31);
r5 = _mm256_permute2f128_ps(tt1, tt5, 0x31);
r6 = _mm256_permute2f128_ps(tt2, tt6, 0x31);
r7 = _mm256_permute2f128_ps(tt3, tt7, 0x31);
}
template<typename T>
inline void transpose_16x16_kernel(float* dst, T* src, size_t dst_stride, size_t src_stride) {
__m256 r0, r1, r2, r3, r4, r5, r6, r7;
for (int i = 0; i < 16; i += 8) {
for (int j = 0; j < 16; j += 8) {
r0 = mm256_uni_loadu_ps(src + src_stride * j);
r1 = mm256_uni_loadu_ps(src + src_stride * (1 + j));
r2 = mm256_uni_loadu_ps(src + src_stride * (2 + j));
r3 = mm256_uni_loadu_ps(src + src_stride * (3 + j));
r4 = mm256_uni_loadu_ps(src + src_stride * (4 + j));
r5 = mm256_uni_loadu_ps(src + src_stride * (5 + j));
r6 = mm256_uni_loadu_ps(src + src_stride * (6 + j));
r7 = mm256_uni_loadu_ps(src + src_stride * (7 + j));
transpose_8x8(r0, r1, r2, r3, r4, r5, r6, r7);
_mm256_storeu_ps(dst + j, r0);
_mm256_storeu_ps(dst + j + dst_stride, r1);
_mm256_storeu_ps(dst + j + dst_stride * 2, r2);
_mm256_storeu_ps(dst + j + dst_stride * 3, r3);
_mm256_storeu_ps(dst + j + dst_stride * 4, r4);
_mm256_storeu_ps(dst + j + dst_stride * 5, r5);
_mm256_storeu_ps(dst + j + dst_stride * 6, r6);
_mm256_storeu_ps(dst + j + dst_stride * 7, r7);
}
src += 8;
dst += 8 * dst_stride;
}
}
#else
template<typename TSRC, typename TDST>
inline void transpose_16x16_kernel(TDST* dst, TSRC* src, size_t dst_stride, size_t src_stride) {
for (size_t i = 0; i < 16; i++) {
for (size_t j = 0; j < 16; j++) {
dst[i * dst_stride + j] = static_cast<TDST>(src[i + j * src_stride]);
}
}
}
#endif
} // namespace XARCH
} // namespace Cpu
} // namespace Extensions
} // namespace ov

41
src/plugins/intel_cpu/src/nodes/kernels/x64/brgemm_kernel.cpp
View File

@ -24,7 +24,8 @@ BrgemmKernel::BrgemmKernel(size_t M,
size_t ldb,
size_t ldc,
bool b_transposed,
ov::element::Type inType)
ov::element::Type inType,
bool b_accumulate)
: M(M),
K(K),
N(N),
@ -32,7 +33,8 @@ BrgemmKernel::BrgemmKernel(size_t M,
ldb(ldb),
ldc(ldc),
b_transposed(b_transposed),
inType(inType) {
inType(inType),
b_accumulate(b_accumulate) {
// blocking M
M_blk = matmulOptimalM;
M_tail = M % M_blk;
@ -45,7 +47,11 @@ BrgemmKernel::BrgemmKernel(size_t M,
THROW_ERROR("brgemm bf16 kernel could only be used above avx512_bf16");
bool isAMXSupported = is_bf16 && mayiuse(avx512_core_amx);
size_t vlen = cpu_isa_traits<avx512_core>::vlen;
size_t vlen;
if (mayiuse(avx512_core))
vlen = cpu_isa_traits<avx512_core>::vlen;
else
vlen = cpu_isa_traits<cpu_isa_t::avx2>::vlen;
// blocking N
N_blk = is_bf16 ? 32 : std::max(N, vlen / inType.size());
N_tail = N % N_blk;
@ -67,7 +73,7 @@ BrgemmKernel::BrgemmKernel(size_t M,
auto M_ = m ? M_tail : M < M_blk ? 0 : M_blk;
auto N_ = n ? N_tail : N - N_tail;
auto K_ = k ? K_tail : K - K % K_blk;
auto beta = k && brgCtxs[getBrgIdx(m, 0, n)].K != 0 ? 1.0f : 0.0f;
auto beta = (b_accumulate || (k && brgCtxs[getBrgIdx(m, 0, n)].K != 0)) ? 1.0f : 0.0f;
brgemmCtx.M = M_;
brgemmCtx.N = N_;
@ -134,9 +140,14 @@ void BrgemmKernel::init_brgemm(brgemmCtx& ctx,
const bool is_int8 =
one_of(ctx.dt_in0, data_type::u8, data_type::s8) && one_of(ctx.dt_in1, data_type::u8, data_type::s8);
auto isa = use_amx ? isa_undef
: ctx.dt_in0 == dnnl_data_type_t::dnnl_bf16 ? avx512_core_bf16
: (is_int8 ? avx512_core_vnni : avx512_core);
cpu_isa_t isa;
if (mayiuse(avx512_core)) {
isa = use_amx ? isa_undef
: ctx.dt_in0 == dnnl_data_type_t::dnnl_bf16 ? avx512_core_bf16
: (is_int8 ? avx512_core_vnni : avx512_core);
} else {
isa = cpu_isa_t::avx2;
}
auto status = brgemm_desc_init(&brgDesc,
isa,
brgemm_addr,
@ -158,6 +169,20 @@ void BrgemmKernel::init_brgemm(brgemmCtx& ctx,
THROW_ERROR("cannot be executed due to invalid brgconv params");
}
if (use_amx && b_accumulate) {
brgemm_attr_t brgattr;
brgattr.max_bs = 1;
brgattr.wary_tail_read = false;
brgattr.hint_innermost_loop = brgemm_innermost_undef;
// if b_accumulate is true, it means we want c+=a*b. jit_brgemm_amx_uker_base_t::load_accumulators can support this using tileload(c) without postops
brgattr.use_uker = true;
brgattr.use_interleave_stores = true;
brgattr.hint_prefetching = brgemm_kernel_prefetching_t::brgemm_prf1;
if (brgemm_desc_set_attr(&brgDesc, brgattr) != dnnl_success) {
THROW_ERROR("cannot be executed due to brgemm_desc_set_attr failed");
}
}
ctx.is_with_amx = use_amx;
status = brgemm_init_tiles(brgDesc, ctx.palette);
if (use_amx) {
@ -319,7 +344,7 @@ void BrgemmKernel::executeGemm(bool is_M_tail, void* a, void* b, void* c, void*
for (size_t k = 0; k < 2; k++) {
size_t mIdx = is_M_tail ? 1 : 0;
auto& brgemmCtx = brgCtxs[getBrgIdx(mIdx, k, n)];
if (brgemmCtx.K != 0 && brgemmCtx.N != 0) {
if (brgemmCtx.K != 0 && brgemmCtx.N != 0 && brgemmCtx.M != 0) {
auto local_a_ptr = k > 0 ? ptr_a_tail : ptr_A;
auto B_stride = (k * count_K + n * count_N * brgVnniFactor) * inType.size();
auto weight_ptr = ptr_scartch_b + B_stride;

4
src/plugins/intel_cpu/src/nodes/kernels/x64/brgemm_kernel.hpp
View File

@ -28,7 +28,8 @@ public:
size_t ldb,
size_t ldc,
bool b_transposed = false,
ov::element::Type inType = ov::element::bf16);
ov::element::Type inType = ov::element::bf16,
bool b_accumulate = false);
// execute all M
void executeGemm(void* a, void* b, void* c, void* wsp, void* scratch_a, void* scratch_b);
// execute by m_blk
@ -58,6 +59,7 @@ private:
size_t packedBSize = 0;
size_t packedASize = 0;
ov::element::Type inType;
bool b_accumulate = false;
static constexpr size_t MHA_BRGEMM_KERNELS_NUM = 8;
static constexpr size_t matmulOptimalM = 32;
struct brgemmCtx {

216
src/plugins/intel_cpu/src/nodes/paged_attn.cpp Normal file
View File

@ -0,0 +1,216 @@
// Copyright (C) 2018-2024 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//
#include "paged_attn.h"
#include "common/arbitrary_order_desc_creator.h"
#include "common/primitive_hashing_utils.hpp"
#include "cpu/x64/cpu_isa_traits.hpp"
#include "dnnl_extension_utils.h"
#include "memory_desc/cpu_memory_desc_utils.h"
#include "memory_desc/dnnl_blocked_memory_desc.h"
#include "onednn/dnnl.h"
#include "openvino/core/parallel.hpp"
#include "openvino/util/common_util.hpp"
#include "shape_inference/custom/paged_attn.hpp"
#include "shape_inference/shape_inference_internal_dyn.hpp"
#include "utils/plain_tensor.hpp"
#include "kernels/scaled_attn/executor_pa.hpp"
#include "kernels/scaled_attn/attn_memcpy.hpp"
#include "kernels/scaled_attn/attn_quant.hpp"
#include <algorithm>
#include <string>
#include <vector>
using namespace ov::Extensions::Cpu;
using namespace ov::Extensions::Cpu::XARCH;
using namespace dnnl::impl;
using namespace dnnl::impl::cpu::x64;
namespace ov {
namespace intel_cpu {
namespace node {
struct PagedAttentionKey {
ov::element::Type rtPrecision;
size_t hash() const;
bool operator==(const PagedAttentionKey& rhs) const;
};
size_t PagedAttentionKey::hash() const {
size_t seed = 0;
seed = hash_combine(seed, rtPrecision.hash());
return seed;
}
bool PagedAttentionKey::operator==(const PagedAttentionKey& rhs) const {
auto retVal = rtPrecision == rhs.rtPrecision;
return retVal;
}
PagedAttention::PagedAttention(const std::shared_ptr<ov::Node>& op, const GraphContext::CPtr context)
: Node(op, context, PAShapeInferFactory(op)) {
std::string errorMessage;
if (!isSupportedOperation(op, errorMessage)) {
OPENVINO_THROW("CPU: " + errorMessage);
}
}
void PagedAttention::initSupportedPrimitiveDescriptors() {
if (!supportedPrimitiveDescriptors.empty())
return;
auto rtPrecision = getRuntimePrecision();
NodeConfig config;
auto& creatorsMap = BlockedDescCreator::getCommonCreators();
auto orgInputNumber = getOriginalInputsNumber();
config.inConfs.resize(orgInputNumber);
config.outConfs.resize(getOriginalOutputsNumber());
config.inConfs[0].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
rtPrecision, getInputShapeAtPort(0)));
config.inConfs[1].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
rtPrecision, getInputShapeAtPort(1)));
config.inConfs[2].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
rtPrecision, getInputShapeAtPort(2)));
OPENVINO_ASSERT(orgInputNumber == 13 || orgInputNumber == 14, "The input number of PagedAttention should be 13 or 14.");
// kvcache, float, []
auto past_kv_input_mem_precision = getOriginalInputPrecisionAtPort(PagedAttentionExecutor::ID_KCACHE);
config.inConfs[PagedAttentionExecutor::ID_KCACHE].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
past_kv_input_mem_precision, getInputShapeAtPort(PagedAttentionExecutor::ID_KCACHE)));
config.inConfs[PagedAttentionExecutor::ID_VCACHE].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
past_kv_input_mem_precision, getInputShapeAtPort(PagedAttentionExecutor::ID_VCACHE)));
// is_prompt, bool, []
config.inConfs[PagedAttentionExecutor::ID_IS_PROMPT].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
ov::element::u8, getInputShapeAtPort(PagedAttentionExecutor::ID_IS_PROMPT)));
// slot_mapping, int, [batch_size, max_context_len]
config.inConfs[PagedAttentionExecutor::ID_SLOT_MAPPING].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
ov::element::i32, getInputShapeAtPort(PagedAttentionExecutor::ID_SLOT_MAPPING)));
// max_context_len, int, []
config.inConfs[PagedAttentionExecutor::ID_MAX_CONTEXT_LEN].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
ov::element::i32, getInputShapeAtPort(PagedAttentionExecutor::ID_MAX_CONTEXT_LEN)));
// context_lens, int, [batch_size]
config.inConfs[PagedAttentionExecutor::ID_CONTEXT_LENS].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
ov::element::i32, getInputShapeAtPort(PagedAttentionExecutor::ID_CONTEXT_LENS)));
// block_tables, int, [batch_size, max_block_per_request]
config.inConfs[PagedAttentionExecutor::ID_BLOCK_TABLES].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
ov::element::i32, getInputShapeAtPort(PagedAttentionExecutor::ID_BLOCK_TABLES)));
// scale, float, []
config.inConfs[PagedAttentionExecutor::ID_SCALE].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
ov::element::f32, getInputShapeAtPort(PagedAttentionExecutor::ID_SCALE)));
// alibi_slopes, float, [?] or nullptr
config.inConfs[PagedAttentionExecutor::ID_ALIBI_SLOPES].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
ov::element::f32, getInputShapeAtPort(PagedAttentionExecutor::ID_ALIBI_SLOPES)));
// sliding_window, int, []
config.inConfs[PagedAttentionExecutor::ID_SLIDING_WINDOW].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
ov::element::i32, getInputShapeAtPort(PagedAttentionExecutor::ID_SLIDING_WINDOW)));
if (orgInputNumber == 14) {
// subsequence_lens, int, [batch_size]
config.inConfs[PagedAttentionExecutor::ID_SUBSEQUENCE_LENS].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
ov::element::i32, getInputShapeAtPort(PagedAttentionExecutor::ID_SUBSEQUENCE_LENS)));
}
config.outConfs[0].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
rtPrecision, getOutputShapeAtPort(0)));
supportedPrimitiveDescriptors.emplace_back(config, impl_desc_type::ref_any);
}
void PagedAttention::createPrimitive() {
auto rtPrecision = getRuntimePrecision();
// in one model, kvCachePrecision could not be changed so no need to care whether it may be changed.
PagedAttentionKey key = {rtPrecision};
auto builder = [&](const PagedAttentionKey& key) -> std::shared_ptr<PagedAttentionExecutor> {
#ifdef OPENVINO_ARCH_X86_64
auto kvCachePrecision = getOriginalInputPrecisionAtPort(PagedAttentionExecutor::ID_KCACHE);
return make_pa_executor(rtPrecision, kvCachePrecision);
#else
return nullptr;
#endif
};
auto cache = context->getParamsCache();
auto result = cache->getOrCreate(key, builder);
if (!result.first) {
OPENVINO_THROW("PagedAttention AttentionExecutor creation fails with precision " + rtPrecision.to_string());
}
m_executor = result.first;
}
void PagedAttention::execute(dnnl::stream strm) {
auto orginInputNumber = getOriginalInputsNumber();
std::vector<MemoryPtr> inputs(orginInputNumber);
auto output = getDstMemoryAtPort(0);
for (size_t i = 0; i < orginInputNumber; i++) {
inputs[i] = getSrcMemoryAtPort(i);
}
gatherConcatPastkvForPagedAttn(inputs);
m_executor->execute(inputs, output);
}
bool PagedAttention::isSupportedOperation(const std::shared_ptr<const ov::Node>& op, std::string& errorMessage) noexcept {
try {
int orgInput = static_cast<int>(op->get_input_size());
if (op->get_type_name() == std::string("PagedAttentionExtension") && orgInput == PagedAttentionExecutor::ID_SLIDING_WINDOW + 1) {
return true;
}
} catch (...) {
return false;
}
return true;
}
void PagedAttention::gatherConcatPastkvForPagedAttn(const std::vector<MemoryPtr>& inputs) {
PlainTensor k, v, k_cache, v_cache, slot_mapping;
k.reset(inputs[PagedAttentionExecutor::ID_K]); // [B, L1, H * S]
v.reset(inputs[PagedAttentionExecutor::ID_V]);
k_cache.reset(inputs[PagedAttentionExecutor::ID_KCACHE]); // [NUM_BLOCKS, H, 32, S]
v_cache.reset(inputs[PagedAttentionExecutor::ID_VCACHE]); // [NUM_BLOCKS, H, 32, S]
slot_mapping.reset(inputs[PagedAttentionExecutor::ID_SLOT_MAPPING]); // [B, max_context_len]
auto B = k.size(0);
auto L1 = k.size(1);
auto H = k_cache.size(1);
auto S = v_cache.size(3) - (k_cache.m_dt == ov::element::Type_t::u8 ? 8 : 0);
k.assert_dims({B, L1, H * S});
v.assert_dims({B, L1, H * S});
slot_mapping.assert_dims({B, 0}, true);
k = k.reshape({B, L1, H, S}).permute({0, 2, 1, 3});
v = v.reshape({B, L1, H, S}).permute({0, 2, 1, 3});
if (k_cache.m_dt == ov::element::Type_t::u8) {
k_cache.assert_dims({0, H, 0, S + 8}, true);
v_cache.assert_dims({k_cache.m_dims[0], H, k_cache.m_dims[2], S + 8});
paged_attn_quantkv(k, v, k_cache, v_cache, slot_mapping);
} else {
k_cache.assert_dims({0, H, 0, S}, true);
v_cache.assert_dims({k_cache.m_dims[0], H, k_cache.m_dims[2], S});
paged_attn_memcpy(k, v, k_cache, v_cache, slot_mapping);
}
}
ov::element::Type PagedAttention::getRuntimePrecision() const {
auto rtPrecision = getOriginalInputPrecisionAtPort(0);
// bf16 should be enabled only when platform supports
if (rtPrecision == ov::element::bf16 && ov::with_cpu_x86_bfloat16()) {
rtPrecision = ov::element::bf16;
} else {
rtPrecision = ov::element::f32;
}
return rtPrecision;
}
} // namespace node
} // namespace intel_cpu
} // namespace ov

51
src/plugins/intel_cpu/src/nodes/paged_attn.h Normal file
View File

@ -0,0 +1,51 @@
// Copyright (C) 2018-2024 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//
#pragma once
#include "memory_state.h"
#include "node.h"
#include "transformations/cpu_opset/common/op/sdpa.hpp"
#include "utils/plain_tensor.hpp"
#include "kernels/scaled_attn/executor_pa.hpp"
namespace ov {
namespace intel_cpu {
namespace node {
class PagedAttention : public Node {
public:
PagedAttention(const std::shared_ptr<ov::Node>& op, const GraphContext::CPtr context);
void getSupportedDescriptors() override {}
bool created() const override {
return getType() == Type::PagedAttention;
}
// pastkv may have zero dimension
bool isExecutable() const override {
return !isInputTensorAtPortEmpty(0) && !isInputTensorAtPortEmpty(1) && !isInputTensorAtPortEmpty(2);
}
bool needPrepareParams() const override {
return false;
}
void executeDynamicImpl(dnnl::stream strm) override {
execute(strm);
}
void initSupportedPrimitiveDescriptors() override;
void execute(dnnl::stream strm) override;
void createPrimitive() override;
static bool isSupportedOperation(const std::shared_ptr<const ov::Node>& op, std::string& errorMessage) noexcept;
private:
void gatherConcatPastkvForPagedAttn(const std::vector<MemoryPtr>& inputs);
ov::element::Type getRuntimePrecision() const override;
std::shared_ptr<ov::Extensions::Cpu::PagedAttentionExecutor> m_executor;
template <typename T> struct AttentionExecutor;
friend struct PagedAttentionKey;
};
} // namespace node
} // namespace intel_cpu
} // namespace ov

652
src/plugins/intel_cpu/src/nodes/scaled_attn.cpp
View File

@ -63,141 +63,6 @@ bool ScaledDotProductAttentionKey::operator==(const ScaledDotProductAttentionKey
return retVal;
}
#ifdef OPENVINO_ARCH_X86_64
// w = query * Key
//
// query: [1, head_size]
// Key : [block_size, head_size]
// w : [1, block_size]
//
// head_size is known at compile time
struct TileConfig {
uint8_t palette_id;
uint8_t startRow;
uint8_t reserved[14];
uint16_t cols[16];
uint8_t rows[16];
void reset(int palette, int _startRow, const std::vector<std::pair<int, int>>& _rows_columnsBytes) {
palette_id = palette;
startRow = _startRow;
unsigned long i;
for (i = 0; i < 14; i++) {
reserved[i] = 0;
}
for (i = 0; i < _rows_columnsBytes.size(); i++) {
rows[i] = _rows_columnsBytes[i].first;
cols[i] = _rows_columnsBytes[i].second;
}
for (; i < 16; i++) {
cols[i] = 0;
rows[i] = 0;
}
}
};
class TileConfiger : public jit_generator {
public:
DECLARE_CPU_JIT_AUX_FUNCTIONS(TileConfiger)
TileConfiger() : jit_generator(jit_name()) {
create_kernel();
}
void generate() override {
Xbyak::Label release;
test(abi_param1, abi_param1);
jz(release);
ldtilecfg(ptr[abi_param1]);
ret();
L(release);
tilerelease();
ret();
}
};
class JitMatMulVecAMX : public jit_generator {
void operator=(const JitMatMulVecAMX&);
public:
DECLARE_CPU_JIT_AUX_FUNCTIONS(JitMatMulVecAMX)
int m_head_size;
int m_block_size;
TileConfiger m_tile_configer;
TileConfig m_tile_cfg;
JitMatMulVecAMX(int head_size, int block_size) : jit_generator(jit_name()), m_head_size(head_size), m_block_size(block_size) {
create_kernel();
m_tile_cfg.reset(1,
0,
{
{16, 4}, // C:0 M x 1 (4b)
{16, 64}, // A:1 M x 32/64 (64b)
{16, 4}, // B:2 32/64 x 1 (4b)
{16, 4}, // B:3
{16, 4}, // B:4
{16, 4}, // B:5
{16, 4}, // B:6
{16, 4}, // B:7
});
}
void tile_config() {
m_tile_configer(&m_tile_cfg);
}
void tile_release() {
m_tile_configer(nullptr);
}
// to save push/pop: do not use `abi_save_gpr_regs`
Xbyak::Reg64 reg_q_addr = abi_param1;
Xbyak::Reg64 reg_k_addr = abi_param2;
Xbyak::Reg64 reg_dst_addr = abi_param3;
Xbyak::Reg64 reg_stride_A = rax;
Xbyak::Reg64 reg_stride_BC = r9;
Xbyak::Tmm tmmC = tmm0;
Xbyak::Tmm tmmA = tmm1;
Xbyak::Tmm tmmB0 = tmm2;
Xbyak::Tmm tmmB1 = tmm3;
Xbyak::Tmm tmmB2 = tmm4;
Xbyak::Tmm tmmB3 = tmm5;
Xbyak::Tmm tmmB4 = tmm6;
Xbyak::Tmm tmmB5 = tmm7;
void generate() override {
mov(reg_stride_A, m_head_size * 2);
mov(reg_stride_BC, 4);
const int kStep = 32;
if ((m_head_size % 32) != 0)
throw std::runtime_error("head size is not multiple of 32");
if ((m_block_size % 16) != 0)
throw std::runtime_error("block size is not multiple of 16");
auto num_B_tiles = m_head_size / kStep;
if (num_B_tiles > 6)
throw std::runtime_error("number of B tiles is bigger than 6");
/*
B(query) head_size x 1
A(key) matrix : block_size x head_size C(dst) block_size x 1
*/
// load query into B tiles
for (int i = 0; i < num_B_tiles; i++) {
tileloadd(Xbyak::Tmm(tmmB0.getIdx() + i), ptr[reg_q_addr + reg_stride_BC + i * 64]);
}
for (int m = 0; m < m_block_size; m += 16) {
tilezero(tmmC);
for (int i = 0; i < num_B_tiles; i++) {
tileloadd(tmmA, ptr[reg_k_addr + reg_stride_A + i * 64]);
tdpbf16ps(tmmC, tmmA, Xbyak::Tmm(tmmB0.getIdx() + i));
}
tilestored(ptr[reg_dst_addr + reg_stride_BC + m * sizeof(float)], tmmC);
add(reg_k_addr, m_head_size * 2 * 16);
}
ret();
}
};
#endif
// default implementation: reference
template <ScaledDotProductAttention::KernelTypes KType, typename T>
struct MHAKernel {
@ -260,8 +125,7 @@ struct MHAKernel {
PlainTensor& output_emb,
bool has_out_transpose,
bool auto_causal,
float d_scale = 0.0f,
size_t sliding_window = 0) {
float d_scale = 0.0f) {
auto B = query.size(0);
auto H = query.size(1);
auto q_len = query.size(2);
@ -313,18 +177,7 @@ struct MHAKernel {
}
// softmax
if (sliding_window) {
size_t start_idx = 0;
auto new_causal = ncausal;
if (ncausal > sliding_window) {
start_idx = ncausal - static_cast<size_t>(sliding_window);
new_causal = sliding_window;
}
softmax(&attn_score[start_idx], new_causal);
memset(&attn_score[0], 0, sizeof(float) * start_idx);
} else {
softmax(&attn_score[0], ncausal);
}
softmax(&attn_score[0], ncausal);
// linearly combine value
word_vec.assign(head_size, 0.0f);
@ -497,8 +350,7 @@ struct MHAKernel<ScaledDotProductAttention::KT_ONEDNN, T> {
PlainTensor& output_emb,
bool has_out_transpose,
bool auto_causal,
float d_scale = 0.0f,
size_t sliding_window = 0) {
float d_scale = 0.0f) {
const auto B = query.size(0);
const auto H = query.size(1);
const auto q_len = query.size(2);
@ -561,39 +413,17 @@ struct MHAKernel<ScaledDotProductAttention::KT_ONEDNN, T> {
for (size_t m = m_start; m < m_end; m++) {
// apply attention mask & sofmax
auto ncausal = auto_causal ? (kv_len - q_len + m + 1) : kv_len;
if (sliding_window) {
size_t start_idx = 0;
auto new_causal = ncausal;
if (ncausal > sliding_window) {
start_idx = ncausal - static_cast<size_t>(sliding_window);
new_causal = sliding_window;
}
attn_softmax(&score.at<float>({b, h, m, start_idx}),
&weight.at<T>({b, h, m, start_idx}),
d_scale,
alibi_ptr + m * alibi_stride,
attn_mask_ptr + m * attn_mask_stride,
cmask_ptr + m * cmask_stride,
select_nfltmax_at_0,
new_causal,
kv_len - start_idx,
precision_of<T>::value,
precision_of<T>::value);
memset(&weight.at<T>({b, h, m, 0}), 0, sizeof(T) * start_idx);
} else {
attn_softmax(&score.at<float>({b, h, m, 0}),
&weight.at<T>({b, h, m, 0}),
d_scale,
alibi_ptr + m * alibi_stride,
attn_mask_ptr + m * attn_mask_stride,
cmask_ptr + m * cmask_stride,
select_nfltmax_at_0,
ncausal,
kv_len,
precision_of<T>::value,
precision_of<T>::value);
}
attn_softmax(&score.at<float>({b, h, m, 0}),
&weight.at<T>({b, h, m, 0}),
d_scale,
alibi_ptr + m * alibi_stride,
attn_mask_ptr + m * attn_mask_stride,
cmask_ptr + m * cmask_stride,
select_nfltmax_at_0,
ncausal,
kv_len,
precision_of<T>::value,
precision_of<T>::value);
}
T* w_ptr = &weight.at<T>({b, h, m_start, 0});
float* fp32_out_ptr;
@ -657,8 +487,7 @@ struct MHAKernel<ScaledDotProductAttention::KT_ONEDNN, T> {
PlainTensor& output_emb,
bool has_out_transpose,
bool auto_causal,
float d_scale = 0.0f,
size_t sliding_window = 0) {
float d_scale = 0.0f) {
auto head_size = query.size(3);
if (d_scale == 0.0f)
d_scale = 1.0f / sqrt(head_size);
@ -672,8 +501,7 @@ struct MHAKernel<ScaledDotProductAttention::KT_ONEDNN, T> {
output_emb,
has_out_transpose,
auto_causal,
d_scale,
sliding_window);
d_scale);
}
};
@ -715,8 +543,7 @@ struct MHAKernel<ScaledDotProductAttention::KT_MLAS, float> {
PlainTensor& output_emb,
bool has_out_transpose,
bool auto_causal,
float d_scale = 0.0f,
size_t sliding_window = 0) {
float d_scale = 0.0f) {
auto B = query.size(0);
auto H = query.size(1);
auto q_len = query.size(2);
@ -806,39 +633,17 @@ struct MHAKernel<ScaledDotProductAttention::KT_MLAS, float> {
for (size_t m = m_start; m < m_end; m++) {
// apply attention mask & sofmax
auto ncausal = auto_causal ? (kv_len - q_len + m + 1) : kv_len;
if (sliding_window) {
size_t start_idx = 0;
auto new_causal = ncausal;
if (ncausal > sliding_window) {
start_idx = ncausal - static_cast<size_t>(sliding_window);
new_causal = sliding_window;
}
attn_softmax(qk + (m - m_start) * qk_m_stride + start_idx,
qk + (m - m_start) * qk_m_stride + start_idx,
d_scale,
alibi_ptr + m * alibi_stride,
attn_mask_ptr + m * attn_mask_stride,
cmask_ptr + m * cmask_stride,
select_nfltmax_at_0,
new_causal,
kv_len - start_idx,
ov::element::f32,
ov::element::f32);
memset(qk + (m - m_start) * qk_m_stride, 0, sizeof(float) * start_idx);
} else {
attn_softmax(qk + (m - m_start) * qk_m_stride,
qk + (m - m_start) * qk_m_stride,
d_scale,
alibi_ptr + m * alibi_stride,
attn_mask_ptr + m * attn_mask_stride,
cmask_ptr + m * cmask_stride,
select_nfltmax_at_0,
ncausal,
kv_len,
ov::element::f32,
ov::element::f32);
}
attn_softmax(qk + (m - m_start) * qk_m_stride,
qk + (m - m_start) * qk_m_stride,
d_scale,
alibi_ptr + m * alibi_stride,
attn_mask_ptr + m * attn_mask_stride,
cmask_ptr + m * cmask_stride,
select_nfltmax_at_0,
ncausal,
kv_len,
ov::element::f32,
ov::element::f32);
}
mlas_sgemm("N",
"N",
@ -864,9 +669,7 @@ struct MHASingleToken {
PlainTensor m_attn_w;
PlainTensor m_temp;
PlainTensor m_head_sum;
#ifdef OPENVINO_ARCH_X86_64
std::shared_ptr<JitMatMulVecAMX> m_gemv;
#endif
MHASingleToken() {}
// Q, K, V is ready, do attention
@ -883,8 +686,6 @@ struct MHASingleToken {
const PlainTensor& attention_mask,
PlainTensor& output_emb,
const PlainTensor& beams,
size_t max_context_len,
const PlainTensor& context_lens,
bool has_out_transpose,
bool auto_causal,
float d_scale,
@ -893,81 +694,13 @@ struct MHASingleToken {
auto B = query.size(0);
auto H = query.size(1);
auto q_len = query.size(2);
bool is_pagedattn = context_lens;
size_t kv_len;
if (is_pagedattn) {
kv_len = max_context_len;
} else {
kv_len = present_key.size(2);
}
kv_len = present_key.size(2);
bool fastpath_valid = false;
#ifdef OPENVINO_ARCH_X86_64
if (is_pagedattn) {
auto S = query.size(3);
size_t block_size = present_value.size(2);
fastpath_valid = mayiuse(amx_bf16) && (S % 32 == 0) && (block_size % 16 == 0) && (S <= 32 * 6) && present_key.get_precision() == ov::element::bf16;
if (fastpath_valid) {
m_attn_w.resize<float>({B, H, q_len, (kv_len + block_size - 1) / block_size * block_size});
if (!m_gemv)
m_gemv = std::make_shared<JitMatMulVecAMX>(static_cast<int>(S), static_cast<int>(block_size));
auto h_group_num = present_value.size(1);
size_t h_each_group_len = 1;
if (h_group_num != H) {
h_each_group_len = H / h_group_num;
}
auto kv_len_in_blocks = beams.m_dims[1];
auto nthr = static_cast<size_t>(parallel_get_max_threads());
size_t real_len = 0;
for (size_t b = 0; b < B; b++)
real_len += static_cast<size_t>(context_lens.ptr<int32_t>()[b]) / block_size;
if (real_len > nthr) {
parallel_for2d_dynamic(B, kv_len_in_blocks, [&](size_t b, size_t pk_in_blocks) {
auto context_len = static_cast<size_t>(context_lens.ptr<int32_t>()[b]);
// kv_len must be valid
auto pk = pk_in_blocks * block_size;
if (pk < context_len) {
m_gemv->tile_config();
auto block_number = beams.ptr<int32_t>(b)[pk_in_blocks];
for (size_t h_group = 0; h_group < h_group_num; h_group++) {
for (size_t pq = 0; pq < q_len; pq++) {
for (size_t h = h_group * h_each_group_len; h < (h_group + 1) * h_each_group_len; h++) {
(*m_gemv)(query.ptr<ov::bfloat16>(b, h, pq), present_key.ptr<ov::bfloat16>(block_number, h_group),
m_attn_w.ptr<float>(b, h, pq) + pk);
}
}
}
m_gemv->tile_release();
}
});
} else {
parallel_for3d_dynamic(B, kv_len_in_blocks, h_group_num, [&](size_t b, size_t pk_in_blocks, size_t h_group) {
auto context_len = static_cast<size_t>(context_lens.ptr<int32_t>()[b]);
// kv_len must be valid
auto pk = pk_in_blocks * block_size;
if (pk < context_len) {
m_gemv->tile_config();
auto block_number = beams.ptr<int32_t>(b)[pk_in_blocks];
for (size_t pq = 0; pq < q_len; pq++) {
for (size_t h = h_group * h_each_group_len; h < (h_group + 1) * h_each_group_len; h++) {
(*m_gemv)(query.ptr<ov::bfloat16>(b, h, pq), present_key.ptr<ov::bfloat16>(block_number, h_group),
m_attn_w.ptr<float>(b, h, pq) + pk);
}
}
m_gemv->tile_release();
}
});
}
}
}
#endif
if (!fastpath_valid) {
// aligned to cache line (64bytes=16*sizeof(float)) to avoid false sharing
m_attn_w.resize<float>({B, H, q_len, (kv_len + 15) / 16 * 16});
}
mha_single_token(query, fastpath_valid ? PlainTensor() : present_key, present_value, alibi_mask, attention_mask, beams, max_context_len,
context_lens, output_emb, m_attn_w, m_temp, has_out_transpose, auto_causal, d_scale, k_scale_zp, v_scale_zp, m_head_sum);
// aligned to cache line (64bytes=16*sizeof(float)) to avoid false sharing
m_attn_w.resize<float>({B, H, q_len, (kv_len + 15) / 16 * 16});
mha_single_token(query, present_key, present_value, alibi_mask, attention_mask, beams,
output_emb, m_attn_w, m_temp, has_out_transpose, auto_causal, d_scale, k_scale_zp, v_scale_zp, m_head_sum);
}
};
@ -995,111 +728,66 @@ struct ScaledDotProductAttention::AttentionExecutor : public ScaledDotProductAtt
bool fuse_causal_attn = config.config.fuse_causal_attn;
bool is_causal = config.config.is_causal;
bool fuse_concat = config.config.fuse_concat;
bool is_pagedattn = config.is_pageattn;
auto input_num = inputs.size();
bool is_prompt = false;
PlainTensor present_key, present_value;
PlainTensor q_input; // f32[B, H, L1, S]
PlainTensor k_input; // f32[B, H|1, L1, S] / [B, H|1, L0+L1, S]
PlainTensor v_input; // f32[B, H|1, L1, S] / [B, H|1, L0+L1, S]
PlainTensor beam_table; // i32[B, max_kvLen]
PlainTensor context_lens;
PlainTensor attn_mask;
PlainTensor output_emb(output);
float scale_input = 0.0f;
size_t B, L1, L0, S;
size_t sliding_window = 0;
size_t max_context_len = 0;
q_input.reset(inputs[0]);
k_input.reset(inputs[1]);
v_input.reset(inputs[2]);
present_key.reset(presentk_input);
present_value.reset(presentv_input);
if (is_pagedattn) {
is_prompt = *inputs[ID_IS_PROMPT]->getDataAs<uint8_t>() == 1;
max_context_len = static_cast<size_t>(*inputs[ID_MAX_CONTEXT_LEN]->getDataAs<int32_t>());
context_lens.reset(inputs[ID_CONTEXT_LENS]);
beam_table.reset(inputs[ID_BLOCK_TABLES]);
scale_input = *inputs[ID_SCALE]->getDataAs<float>();
// TODO: alibi and sliding window
// no attn mask, auto-generated casual mask
is_causal = true;
has_out_transpose = true;
// q: [B, L1, H*S], kv: [B, L1, Hk*S]
// k_cache: [NUM_BLOCKS, Hk, 16, S]
// v_cache: [NUM_BLOCKS, Hk, 16, S]
// context_lens: [B]
// block_tables: [B, max_block_per_request]
B = k_input.size(0);
L1 = k_input.size(1);
auto Hk = present_key.size(1);
// The layout for per token per head for u8 kv cache:
// |scale(f32)|zeropoint(f32)|quantized feature(u8,idx_1)|quantized feature(u8,idx_2)|...|quantized feature(u8,idx_S)|
// The actual size needs to deduct scale and zeropoint.
S = present_value.size(3) - (present_value.m_dt == ov::element::Type_t::u8 ? sizeof(float) * 2 : 0);
auto H = q_input.size(2) / S;
// L0 in each batch may be different
L0 = 0;
q_input.assert_dims({B, L1, H * S});
if (!is_prompt) {
context_lens.assert_dims({B});
beam_table.assert_dims({B, 0}, true);
if (beam_input)
beam_table.reset(beam_input);
if (input_num > 3) {
// attn_mask
if (inputs[3]->getDesc().getPrecision() == ov::element::u8) {
// bool->f32
prepare_attn_mask(inputs[3]);
attn_mask = attn_buf;
} else {
sliding_window = static_cast<size_t>(*inputs[ID_SLIDING_WINDOW]->getDataAs<int32_t>());
attn_mask.reset(inputs[3]);
}
output_emb.assert_dims({B, L1, H * S});
q_input = q_input.reshape({B, L1, H, S}).permute({0, 2, 1, 3});
k_input = k_input.reshape({B, L1, Hk, S}).permute({0, 2, 1, 3});
v_input = v_input.reshape({B, L1, Hk, S}).permute({0, 2, 1, 3});
} else {
if (beam_input)
beam_table.reset(beam_input);
if (input_num > 3) {
// attn_mask
if (inputs[3]->getDesc().getPrecision() == ov::element::u8) {
// bool->f32
prepare_attn_mask(inputs[3]);
attn_mask = attn_buf;
} else {
attn_mask.reset(inputs[3]);
}
// if has scale, attn_mask must be present
if (input_num > 4) {
scale_input = *inputs[4]->getDataAs<float>();
}
// if has scale, attn_mask must be present
if (input_num > 4) {
scale_input = *inputs[4]->getDataAs<float>();
}
// q: [B, H, L1, S]
const auto & permute_axes = config.config.permute_axes;
if (!permute_axes.empty()) {
q_input = q_input.permute(permute_axes);
k_input = k_input.permute(permute_axes);
v_input = v_input.permute(permute_axes);
present_key = present_key.permute(permute_axes);
present_value = present_value.permute(permute_axes);
}
B = q_input.size(0);
L1 = q_input.size(2);
S = q_input.size(3);
L0 = present_key.size(2) - L1;
auto Hk = k_input.size(1);
if (fuse_concat) {
k_input.assert_dims({B, Hk, L1, S});
v_input.assert_dims({B, Hk, L1, S});
} else {
k_input.assert_dims({B, Hk, L0 + L1, S});
v_input.assert_dims({B, Hk, L0 + L1, S});
}
present_key.assert_dims({B, Hk, L0 + L1, S});
present_value.assert_dims({B, Hk, L0 + L1, S});
if (beam_table)
beam_table.assert_dims({B, L0 + L1});
}
// q: [B, H, L1, S]
const auto & permute_axes = config.config.permute_axes;
if (!permute_axes.empty()) {
q_input = q_input.permute(permute_axes);
k_input = k_input.permute(permute_axes);
v_input = v_input.permute(permute_axes);
present_key = present_key.permute(permute_axes);
present_value = present_value.permute(permute_axes);
}
B = q_input.size(0);
L1 = q_input.size(2);
S = q_input.size(3);
L0 = present_key.size(2) - L1;
auto Hk = k_input.size(1);
if (fuse_concat) {
k_input.assert_dims({B, Hk, L1, S});
v_input.assert_dims({B, Hk, L1, S});
} else {
k_input.assert_dims({B, Hk, L0 + L1, S});
v_input.assert_dims({B, Hk, L0 + L1, S});
}
present_key.assert_dims({B, Hk, L0 + L1, S});
present_value.assert_dims({B, Hk, L0 + L1, S});
if (beam_table)
beam_table.assert_dims({B, L0 + L1});
bool auto_causal;
bool use_attn_mask;
if (fuse_causal_attn) {
@ -1130,15 +818,11 @@ struct ScaledDotProductAttention::AttentionExecutor : public ScaledDotProductAtt
}
// second token, or first token with pastkv fusing
bool use_one_token;
if (is_pagedattn)
use_one_token = !is_prompt;
else
use_one_token = L1 == 1 || (fuse_concat && L0 > 0);
bool use_one_token = L1 == 1 || (fuse_concat && L0 > 0);
if (!use_one_token) {
// multi-token version
kernel(strm, q_input, k_input, v_input, {}, use_attn_mask ? attn_mask : PlainTensor(),
output_emb, has_out_transpose, auto_causal, scale_input, sliding_window);
output_emb, has_out_transpose, auto_causal, scale_input);
} else {
// 1-token version
// for second token, using a special AVX2/AVX512 float path:
@ -1146,7 +830,7 @@ struct ScaledDotProductAttention::AttentionExecutor : public ScaledDotProductAtt
// 2, using float will save the repack cost which typically is required for bf16/int8 opt
// 3, using dot product can leverage the SIMD while easily adapt to indirect kv cache
kernel_single_token(q_input, present_key, present_value, {}, use_attn_mask ? attn_mask : PlainTensor(),
output_emb, beam_table, max_context_len, context_lens, has_out_transpose, auto_causal, scale_input, k_scale_zp, v_scale_zp);
output_emb, beam_table, has_out_transpose, auto_causal, scale_input, k_scale_zp, v_scale_zp);
}
}
};
@ -1158,18 +842,12 @@ ScaledDotProductAttention::ScaledDotProductAttention(const std::shared_ptr<ov::N
OPENVINO_THROW("CPU: " + errorMessage);
}
if (op->get_type_name() == std::string("PagedAttentionExtension")) {
m_is_pageattn = true;
m_config.is_pageattn = true;
const auto node = std::dynamic_pointer_cast<const ov::op::v13::ScaledDotProductAttention>(op);
if (node) {
m_config.config.is_causal = node->get_causal();
} else {
m_is_pageattn = false;
const auto node = std::dynamic_pointer_cast<const ov::op::v13::ScaledDotProductAttention>(op);
if (node) {
m_config.config.is_causal = node->get_causal();
} else {
const auto node = std::dynamic_pointer_cast<const ScaledDotProductAttentionWithKVCache>(op);
m_config.config = node->get_config();
}
const auto node = std::dynamic_pointer_cast<const ScaledDotProductAttentionWithKVCache>(op);
m_config.config = node->get_config();
}
}
@ -1189,83 +867,49 @@ void ScaledDotProductAttention::initSupportedPrimitiveDescriptors() {
rtPrecision, getInputShapeAtPort(1)));
config.inConfs[2].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
rtPrecision, getInputShapeAtPort(2)));
if (m_is_pageattn) {
OPENVINO_ASSERT(getOriginalInputsNumber() == 13, "The input number of PagedAttention should be 13.");
// kvcache, float, []
auto past_kv_input_mem_precision = getOriginalInputPrecisionAtPort(ID_KCACHE);
config.inConfs[ID_KCACHE].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
past_kv_input_mem_precision, getInputShapeAtPort(ID_KCACHE)));
config.inConfs[ID_VCACHE].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
past_kv_input_mem_precision, getInputShapeAtPort(ID_VCACHE)));
// is_prompt, bool, []
config.inConfs[ID_IS_PROMPT].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
ov::element::u8, getInputShapeAtPort(ID_IS_PROMPT)));
// slot_mapping, int, [batch_size, max_context_len]
config.inConfs[ID_SLOT_MAPPING].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
ov::element::i32, getInputShapeAtPort(ID_SLOT_MAPPING)));
// max_context_len, int, []
config.inConfs[ID_MAX_CONTEXT_LEN].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
ov::element::i32, getInputShapeAtPort(ID_MAX_CONTEXT_LEN)));
// context_lens, int, [batch_size]
config.inConfs[ID_CONTEXT_LENS].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
ov::element::i32, getInputShapeAtPort(ID_CONTEXT_LENS)));
// block_tables, int, [batch_size, max_block_per_request]
config.inConfs[ID_BLOCK_TABLES].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
ov::element::i32, getInputShapeAtPort(ID_BLOCK_TABLES)));
// scale, float, []
config.inConfs[ID_SCALE].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
ov::element::f32, getInputShapeAtPort(ID_SCALE)));
// alibi_slopes, float, [?] or nullptr
config.inConfs[ID_ALIBI_SLOPES].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
ov::element::f32, getInputShapeAtPort(ID_ALIBI_SLOPES)));
// sliding_window, int, []
config.inConfs[ID_SLIDING_WINDOW].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
ov::element::i32, getInputShapeAtPort(ID_SLIDING_WINDOW)));
} else {
auto nextPortIdx = 3;
if (orginSDPInputNumber > 3) {
// attn_mask
if (getOriginalInputPrecisionAtPort(nextPortIdx) == ov::element::u8) {
config.inConfs[nextPortIdx].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
ov::element::u8, getInputShapeAtPort(nextPortIdx)));
} else {
config.inConfs[nextPortIdx].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
rtPrecision, getInputShapeAtPort(nextPortIdx)));
}
nextPortIdx++;
}
if (orginSDPInputNumber > 4) {
auto nextPortIdx = 3;
if (orginSDPInputNumber > 3) {
// attn_mask
if (getOriginalInputPrecisionAtPort(nextPortIdx) == ov::element::u8) {
config.inConfs[nextPortIdx].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
ov::element::f32, getInputShapeAtPort(nextPortIdx)));
ov::element::u8, getInputShapeAtPort(nextPortIdx)));
} else {
config.inConfs[nextPortIdx].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
rtPrecision, getInputShapeAtPort(nextPortIdx)));
}
nextPortIdx++;
}
if (orginSDPInputNumber > 4) {
config.inConfs[nextPortIdx].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
ov::element::f32, getInputShapeAtPort(nextPortIdx)));
}
if (m_config.config.fuse_concat) {
// beam_idx
config.inConfs[orginSDPInputNumber + 0].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
ov::element::i32, getInputShapeAtPort(orginSDPInputNumber + 0)));
if (m_config.config.fuse_concat) {
// beam_idx
config.inConfs[orginSDPInputNumber + 0].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
ov::element::i32, getInputShapeAtPort(orginSDPInputNumber + 0)));
// Since the InputMemory nodes are simple proxy for the state memory as well as the init subgraph memory,
// it doesn't make sense to set the real KV cache precision, since we don't need any precision conversions
// provided by the common graph logic. We set precisions equal to the precisions of the state nodes to avoid
// reorder insertion in between MemoryInputSDPA and SDPA nodes.
// Since the InputMemory nodes are simple proxy for the state memory as well as the init subgraph memory,
// it doesn't make sense to set the real KV cache precision, since we don't need any precision conversions
// provided by the common graph logic. We set precisions equal to the precisions of the state nodes to avoid
// reorder insertion in between MemoryInputSDPA and SDPA nodes.
auto past_k_input_mem_precision = getParentEdgeAt(orginSDPInputNumber + 1)->getParent()->getOriginalOutputPrecisionAtPort(0);
// pastk
config.inConfs[orginSDPInputNumber + 1].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
past_k_input_mem_precision, getInputShapeAtPort(orginSDPInputNumber + 1)));
auto past_k_input_mem_precision = getParentEdgeAt(orginSDPInputNumber + 1)->getParent()->getOriginalOutputPrecisionAtPort(0);
// pastk
config.inConfs[orginSDPInputNumber + 1].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
past_k_input_mem_precision, getInputShapeAtPort(orginSDPInputNumber + 1)));
auto past_v_input_mem_precision = getParentEdgeAt(orginSDPInputNumber + 2)->getParent()->getOriginalOutputPrecisionAtPort(0);
// pastv
config.inConfs[orginSDPInputNumber + 2].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
past_v_input_mem_precision, getInputShapeAtPort(orginSDPInputNumber + 2)));
auto past_v_input_mem_precision = getParentEdgeAt(orginSDPInputNumber + 2)->getParent()->getOriginalOutputPrecisionAtPort(0);
// pastv
config.inConfs[orginSDPInputNumber + 2].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
past_v_input_mem_precision, getInputShapeAtPort(orginSDPInputNumber + 2)));
config.outConfs[1].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
past_k_input_mem_precision, getOutputShapeAtPort(1)));
config.outConfs[1].inPlace(-1);
config.outConfs[2].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
past_v_input_mem_precision, getOutputShapeAtPort(2)));
config.outConfs[2].inPlace(-1);
}
config.outConfs[1].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
past_k_input_mem_precision, getOutputShapeAtPort(1)));
config.outConfs[1].inPlace(-1);
config.outConfs[2].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
past_v_input_mem_precision, getOutputShapeAtPort(2)));
config.outConfs[2].inPlace(-1);
}
config.outConfs[0].setMemDesc(creatorsMap.at(LayoutType::ncsp)->createSharedDesc(
@ -1324,35 +968,25 @@ void ScaledDotProductAttention::execute(dnnl::stream strm) {
}
PlainTensor k_scale_zp, v_scale_zp;
if (m_is_pageattn) {
gatherConcatPastkvForPagedAttn(inputs);
if (m_config.config.fuse_concat) {
CPU_NODE_ASSERT(m_k_state && m_v_state, "has null input states");
// initialization will be also completed in this func
gatherConcatPastkv(inputs[1], inputs[2], getSrcMemoryAtPort(orginSDPInputNumber));
presentk_input = inputs[ID_KCACHE];
presentv_input = inputs[ID_VCACHE];
presentk_input = m_k_state->internal_state_mem();
presentv_input = m_v_state->internal_state_mem();
beam_input = m_k_state->hidden_state_mem();
k_scale_zp = m_k_state->get_scale_zp();
v_scale_zp = m_v_state->get_scale_zp();
} else {
if (m_config.config.fuse_concat) {
CPU_NODE_ASSERT(m_k_state && m_v_state, "has null input states");
// initialization will be also completed in this func
gatherConcatPastkv(inputs[1], inputs[2], getSrcMemoryAtPort(orginSDPInputNumber));
presentk_input = m_k_state->internal_state_mem();
presentv_input = m_v_state->internal_state_mem();
beam_input = m_k_state->hidden_state_mem();
k_scale_zp = m_k_state->get_scale_zp();
v_scale_zp = m_v_state->get_scale_zp();
} else {
presentk_input = inputs[1];
presentv_input = inputs[2];
}
presentk_input = inputs[1];
presentv_input = inputs[2];
}
m_executor->execute(strm, m_config, inputs, output, presentk_input, presentv_input, beam_input, k_scale_zp, v_scale_zp);
}
bool ScaledDotProductAttention::isSupportedOperation(const std::shared_ptr<const ov::Node>& op, std::string& errorMessage) noexcept {
try {
if (op->get_type_name() == std::string("PagedAttentionExtension")) {
return true;
}
if (!std::dynamic_pointer_cast<const ov::op::v13::ScaledDotProductAttention>(op) &&
!std::dynamic_pointer_cast<const ScaledDotProductAttentionWithKVCache>(op)) {
errorMessage = "Only ScaledDotProductAttention or ScaledDotProductAttentionWithKVCache operation are supported";
@ -1575,36 +1209,6 @@ void ScaledDotProductAttention::resetBeamTablePastkv(const MemoryPtr& mem_cur_k,
}
}
void ScaledDotProductAttention::gatherConcatPastkvForPagedAttn(const std::vector<MemoryPtr>& inputs) {
PlainTensor k, v, k_cache, v_cache, slot_mapping;
k.reset(inputs[ID_K]); // [B, L1, H * S]
v.reset(inputs[ID_V]);
k_cache.reset(inputs[ID_KCACHE]); // [NUM_BLOCKS, H, 16, S]
v_cache.reset(inputs[ID_VCACHE]); // [NUM_BLOCKS, H, 16, S]
slot_mapping.reset(inputs[ID_SLOT_MAPPING]); // [B, max_context_len]
auto B = k.size(0);
auto L1 = k.size(1);
auto H = k_cache.size(1);
auto S = v_cache.size(3) - (k_cache.m_dt == ov::element::Type_t::u8 ? 8 : 0);
k.assert_dims({B, L1, H * S});
v.assert_dims({B, L1, H * S});
slot_mapping.assert_dims({B, 0}, true);
k = k.reshape({B, L1, H, S}).permute({0, 2, 1, 3});
v = v.reshape({B, L1, H, S}).permute({0, 2, 1, 3});
if (k_cache.m_dt == ov::element::Type_t::u8) {
k_cache.assert_dims({0, H, 0, S + 8}, true);
v_cache.assert_dims({k_cache.m_dims[0], H, k_cache.m_dims[2], S + 8});
paged_attn_quantkv(k, v, k_cache, v_cache, slot_mapping);
} else {
k_cache.assert_dims({0, H, 0, S}, true);
v_cache.assert_dims({k_cache.m_dims[0], H, k_cache.m_dims[2], S});
paged_attn_memcpy(k, v, k_cache, v_cache, slot_mapping);
}
}
void ScaledDotProductAttention::gatherConcatPastkv(const MemoryPtr& mem_cur_k, const MemoryPtr& mem_cur_v, const MemoryPtr& mem_beam_idx) {
PlainTensor cur_k;
cur_k.reset(mem_cur_k);

18
src/plugins/intel_cpu/src/nodes/scaled_attn.h
View File

@ -52,7 +52,6 @@ public:
private:
void gatherConcatPastkv(const MemoryPtr& mem_cur_k, const MemoryPtr& mem_cur_v, const MemoryPtr& mem_beam_idx);
void gatherConcatPastkvForPagedAttn(const std::vector<MemoryPtr>& inputs);
void updateBeamTable(const MemoryPtr& mem_beam_idx, size_t new_q_len);
void updatePastkv(const MemoryPtr& mem_cur_k, const MemoryPtr& mem_cur_v);
ov::element::Type getRuntimePrecision() const override;
@ -60,7 +59,6 @@ private:
struct Config {
ScaledDotProductAttentionWithKVCache::Config config;
bool is_pageattn = false;
};
struct Executor {
@ -70,7 +68,6 @@ private:
virtual ~Executor() = default;
};
bool m_is_pageattn;
Config m_config;
std::shared_ptr<Executor> m_executor;
template <KernelTypes KType, typename T> struct AttentionExecutor;
@ -82,21 +79,6 @@ private:
// (0, 1, 2, 3) for BHLS
// (2, 0, 1, 3) for LBHS
std::vector<size_t> m_kvstate_layout = {2, 0, 1, 3};
// PagedAttention input index
static const size_t ID_Q = 0;
static const size_t ID_K = 1;
static const size_t ID_V = 2;
static const size_t ID_KCACHE = 3;
static const size_t ID_VCACHE = 4;
static const size_t ID_IS_PROMPT = 5;
static const size_t ID_SLOT_MAPPING = 6;
static const size_t ID_MAX_CONTEXT_LEN = 7;
static const size_t ID_CONTEXT_LENS = 8;
static const size_t ID_BLOCK_TABLES = 9;
static const size_t ID_SCALE = 10;
static const size_t ID_ALIBI_SLOPES = 11;
static const size_t ID_SLIDING_WINDOW = 12;
};
} // namespace node

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

@ -61,6 +61,7 @@
#include "nodes/normalize.h"
#include "nodes/one_hot.h"
#include "nodes/pad.h"
#include "nodes/paged_attn.h"
#include "nodes/pooling.h"
#include "nodes/priorbox.h"
#include "nodes/priorbox_clustered.h"
@ -205,6 +206,7 @@ Node::NodesFactory::NodesFactory() : Factory("NodesFactory") {
INTEL_CPU_NODE(Interaction, Type::Interaction);
INTEL_CPU_NODE(MHA, Type::MHA);
INTEL_CPU_NODE(ScaledDotProductAttention, Type::ScaledDotProductAttention);
INTEL_CPU_NODE(PagedAttention, Type::PagedAttention);
INTEL_CPU_NODE(Snippet, Type::Subgraph);
#endif
}

38
src/plugins/intel_cpu/src/shape_inference/custom/paged_attn.cpp Normal file
View File

@ -0,0 +1,38 @@
// Copyright (C) 2018-2024 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//
#include "paged_attn.hpp"
#include "shape_inference/shape_inference_cpu.hpp"
#include "shape_inference/shape_inference_ngraph.hpp"
#include "transformations/cpu_opset/common/op/sdpa.hpp"
#include "utils.hpp"
namespace ov {
namespace intel_cpu {
namespace node {
class PAShapeInfer : public ShapeInferEmptyPads {
public:
PAShapeInfer() {}
IShapeInfer::Result infer(const std::vector<std::reference_wrapper<const VectorDims>>& input_shapes,
const std::unordered_map<size_t, MemoryPtr>& data_dependency) override {
const auto& query_dims = input_shapes.front().get();
return {{query_dims}, ShapeInferStatus::success};
}
port_mask_t get_port_mask() const override {
return EMPTY_PORT_MASK;
}
};
ShapeInferPtr PAShapeInferFactory::makeShapeInfer() const {
return std::make_shared<PAShapeInfer>();
}
} // namespace node
} // namespace intel_cpu
} // namespace ov

24
src/plugins/intel_cpu/src/shape_inference/custom/paged_attn.hpp Normal file
View File

@ -0,0 +1,24 @@
// Copyright (C) 2018-2024 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//
#include <node.h>
#include "shape_inference/shape_inference_cpu.hpp"
#pragma once
namespace ov {
namespace intel_cpu {
namespace node {
class PAShapeInferFactory : public ShapeInferFactory {
public:
PAShapeInferFactory(std::shared_ptr<ov::Node> op) : m_op(op) {}
ShapeInferPtr makeShapeInfer() const override;
private:
std::shared_ptr<ov::Node> m_op;
};
} // namespace node
} // namespace intel_cpu
} // namespace ov