Added framework map to map to PyTorch module

openvinotoolkit · Feb 5, 2024 · e9d5e44 · e9d5e44
1 parent b6c35d6
commit e9d5e44
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 0 deletions.
diff --git a/modules/custom_operations/user_ie_extensions/CMakeLists.txt b/modules/custom_operations/user_ie_extensions/CMakeLists.txt
@@ -100,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})

diff --git a/modules/custom_operations/user_ie_extensions/ov_extension.cpp b/modules/custom_operations/user_ie_extensions/ov_extension.cpp
@@ -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>>(),
+#else
+#    define PAGED_ATTENTION_EXT
+#endif
+
 OPENVINO_CREATE_EXTENSIONS(std::vector<ov::Extension::Ptr>(
     {
         CALCULATE_GRID_EXT
@@ -94,4 +103,5 @@ OPENVINO_CREATE_EXTENSIONS(std::vector<ov::Extension::Ptr>(
         S_CONV_EXT
         COMPLEX_MUL_EXT
         TOKENIZER_EXT
+        PAGED_ATTENTION_EXT
     }));
diff --git a/modules/custom_operations/user_ie_extensions/paged_attention/paged_attention.hpp b/modules/custom_operations/user_ie_extensions/paged_attention/paged_attention.hpp
@@ -6,6 +6,7 @@
 
 #include "openvino/op/op.hpp"
 #include "openvino/runtime/infer_request.hpp"
+#include "openvino/frontend/pytorch/extension/op.hpp"
 
 namespace TemplateExtension {
 
@@ -23,6 +24,7 @@ namespace TemplateExtension {
 class PagedAttention : public ov::op::Op {
 public:
     OPENVINO_OP("PagedAttention");
+    OPENVINO_FRAMEWORK_MAP(pytorch, "vllm.model_executor.layers.attention.PagedAttention");
 
     PagedAttention() = default;