Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

PagedAttention operation #867

Closed
wants to merge 22 commits into from
Closed

PagedAttention operation #867

Show file tree
Hide file tree
Changes from 8 commits
Commits
Show all changes
22 commits
Select commit Hold shift + click to select a range
dc4553a
Introduced first version of PagedAttention operation
ilya-lavrenov Feb 2, 2024
1542b7a
Added CPU implementation for reshape_and_cache
ilya-lavrenov Feb 2, 2024
cbeff36
Attention can be compiled
ilya-lavrenov Feb 2, 2024
cb708a0
Temp commit
ilya-lavrenov Feb 2, 2024
b340cb3
Creare SPDA operation for prefill case
ilya-lavrenov Feb 5, 2024
b6c35d6
Finalize PagedAttention implementation
ilya-lavrenov Feb 5, 2024
e9d5e44
Added framework map to map to PyTorch module
ilya-lavrenov Feb 5, 2024
04e7014
Use Tensor::set_shape instead of view
ilya-lavrenov Feb 5, 2024
675de7b
Intermediate state
ilya-lavrenov Feb 7, 2024
980fabc
Passes accuracy
ilya-lavrenov Feb 11, 2024
b36e57a
Merge remote-tracking branch 'upstream/master' into paged-attention
ilya-lavrenov Feb 11, 2024
7f9ae41
No explicit registration of frontend::OpExtension, removed unnecessar…
slyalin Feb 13, 2024
e1eda1a
Merge pull request #1 from slyalin/simplified_extension_registration
ilya-lavrenov Feb 14, 2024
3a368f8
Merge remote-tracking branch 'upstream/master' into paged-attention
ilya-lavrenov Feb 15, 2024
5b17873
Remove print of performance
ilya-lavrenov Feb 15, 2024
f4bc7c8
Supported GQA, MQA
ilya-lavrenov Feb 15, 2024
264e86d
microsoft/phi-1_5 works
ilya-lavrenov Feb 17, 2024
793350d
Suppport mistral with sliding_windows
ilya-lavrenov Feb 18, 2024
c4705ac
Fixed sliding_window
ilya-lavrenov Feb 18, 2024
bf920e3
Passing alibi_slopes and sliding_window to PagedAttension kernel
slyalin Feb 19, 2024
4e1cfd3
Fixed overflow with maximal sliding_window
slyalin Feb 19, 2024
d6388a9
Merge pull request #2 from slyalin/paged_attention_remaining_parameters
ilya-lavrenov Feb 19, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Empty file added modules/custom_operations/tests/CMakeLists.txt
View file
Open in desktop
Empty file.
9 changes: 6 additions & 3 deletions modules/custom_operations/user_ie_extensions/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,9 +12,7 @@ endif()

set(TARGET_NAME "user_ov_extensions")

if(NOT CMAKE_CXX_STANDARD)
set(CMAKE_CXX_STANDARD 11)
endif()
set(CMAKE_CXX_STANDARD 17)

find_package(OpenVINO REQUIRED COMPONENTS Runtime)
find_package(TBB COMPONENTS tbb)
Expand Down Expand Up @@ -102,6 +100,11 @@ if(tokenizer IN_LIST CUSTOM_OPERATIONS)
endif()
endif()

if(paged_attention IN_LIST CUSTOM_OPERATIONS)
find_package(OpenVINO REQUIRED COMPONENTS PyTorch)
target_link_libraries(${TARGET_NAME} PRIVATE openvino::frontend::pytorch)
endif()

target_link_libraries(${TARGET_NAME} PRIVATE openvino::runtime)

target_compile_definitions(${TARGET_NAME} PRIVATE IMPLEMENT_OPENVINO_EXTENSION_API ${CUSTOM_OPERATIONS})
Expand Down
10 changes: 10 additions & 0 deletions modules/custom_operations/user_ie_extensions/ov_extension.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,15 @@
# define TOKENIZER_EXT
#endif

#ifdef paged_attention
# include "paged_attention/paged_attention.hpp"
# define PAGED_ATTENTION_EXT \
std::make_shared<ov::OpExtension<TemplateExtension::PagedAttention>>(), \
std::make_shared<ov::frontend::OpExtension<TemplateExtension::PagedAttention>>(),
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

With this line, do we still need openvinotoolkit/openvino#22656?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just checked, openvinotoolkit/openvino#22656 is not needed.

#else
# define PAGED_ATTENTION_EXT
#endif

OPENVINO_CREATE_EXTENSIONS(std::vector<ov::Extension::Ptr>(
{
CALCULATE_GRID_EXT
Expand All @@ -94,4 +103,5 @@ OPENVINO_CREATE_EXTENSIONS(std::vector<ov::Extension::Ptr>(
S_CONV_EXT
COMPLEX_MUL_EXT
TOKENIZER_EXT
PAGED_ATTENTION_EXT
}));
Loading