diff --git a/python/env_dict.py.in b/python/env_dict.py.in
index e5bd6a0619d44e7615d44eb3a62431a6e04cb812..d06a9d1af3a67f131430eff397f77b107a7307d1 100644
--- a/python/env_dict.py.in
+++ b/python/env_dict.py.in
@@ -71,5 +71,6 @@ env_dict={
     'WIN32':'@WIN32@',
     'JIT_RELEASE_WHL':'@JIT_RELEASE_WHL@',
     'WITH_PSLIB':'@WITH_PSLIB@',
-    'PYBIND_INCLUDE_DIR':'@PYBIND_INCLUDE_DIR@'
+    'PYBIND_INCLUDE_DIR':'@PYBIND_INCLUDE_DIR@',
+    'WITH_XDNN_API':'@WITH_XDNN_API@'
 }
diff --git a/python/setup.py.in b/python/setup.py.in
index f92eeb39c1d650aa32eb8b8fb5924a9815580ad5..452ca01e2b6bab8455650777f684fff37c2c6e70 100644
--- a/python/setup.py.in
+++ b/python/setup.py.in
@@ -728,6 +728,9 @@ if '${WITH_GPU}' == 'ON' or '${WITH_ROCM}' == 'ON':
     # externalErrorMsg.pb for External Error message
     headers += list(find_files('*.pb', '${externalError_INCLUDE_DIR}'))
 
+if '${WITH_XDNN_API}' == 'ON':
+    headers += list(find_files('*.h', '@PADDLE_BINARY_DIR@/third_party/xpu/src/extern_xpu/xpu', recursive=True)) # xdnn api headers
+
 headers += list(find_files('*.h', '${PYBIND_INCLUDE_DIR}', True)) # pybind headers
 
 class InstallCommand(InstallCommandBase):
@@ -776,7 +779,7 @@ class InstallHeaders(Command):
         else:
             # third_party
             install_dir = re.sub('${THIRD_PARTY_PATH}', 'third_party', header)
-            patterns = ['install/mkldnn/include', 'pybind/src/extern_pybind/include']
+            patterns = ['install/mkldnn/include', 'pybind/src/extern_pybind/include', 'third_party/xpu/src/extern_xpu/xpu/include/']
             for pattern in patterns:
                 install_dir = re.sub(pattern, '', install_dir)
         install_dir = os.path.join(self.install_dir, os.path.dirname(install_dir))
diff --git a/setup.py b/setup.py
index 4e6c0436f0a211f6ce18da4eaa6339bd0cf89669..07c96b720fa2364ba2bd81d0f7ef7fe82d10b416 100644
--- a/setup.py
+++ b/setup.py
@@ -155,6 +155,7 @@ def get_header_install_dir(header):
         patterns = [
             'install/mkldnn/include',
             'pybind/src/extern_pybind/include',
+            'third_party/xpu/src/extern_xpu/xpu/include/',
         ]
         for pattern in patterns:
             install_dir = re.sub(pattern, '', install_dir)
@@ -1222,6 +1223,15 @@ def get_headers():
             find_files('*.pb', env_dict.get("externalError_INCLUDE_DIR"))
         )
 
+    if env_dict.get("WITH_XDNN_API") == 'ON':
+        headers += list(
+            find_files(
+                '*.h',
+                paddle_binary_dir + '/third_party/xpu/src/extern_xpu/xpu',
+                recursive=True,
+            )
+        )  # xdnn api headers
+
     # pybind headers
     headers += list(find_files('*.h', env_dict.get("PYBIND_INCLUDE_DIR"), True))
     return headers