diff --git a/paddle/pten/CMakeLists.txt b/paddle/pten/CMakeLists.txt
index 7f6a64a1bb2e79eed409f1b43b4cf32a80365863..4538016333e5e76663c328c7b3e78c3ceb2b7fda 100644
--- a/paddle/pten/CMakeLists.txt
+++ b/paddle/pten/CMakeLists.txt
@@ -23,7 +23,7 @@ add_subdirectory(ops)
 add_subdirectory(tests)
 
 # make an unity target for compile deps
-set(PTEN_DEPS convert_utils dense_tensor kernel_factory kernel_context)
+set(PTEN_DEPS convert_utils dense_tensor pten_context kernel_factory kernel_context)
 set(PTEN_DEPS ${PTEN_DEPS} math_cpu linalg_cpu creation_cpu manipulation_cpu)
 set(PTEN_DEPS ${PTEN_DEPS} nary unary binary)
 if(WITH_GPU OR WITH_ROCM)
diff --git a/paddle/pten/api/lib/CMakeLists.txt b/paddle/pten/api/lib/CMakeLists.txt
index d1e60c4505d6b0ad76b695940740c7a03b1ede18..1c2b3823920d6e7877842ab48d66f48f4a9af076 100644
--- a/paddle/pten/api/lib/CMakeLists.txt
+++ b/paddle/pten/api/lib/CMakeLists.txt
@@ -10,7 +10,7 @@ else()
   cc_library(pten_tensor SRCS tensor.cc DEPS tensor_base dense_tensor pten_api_utils ext_compat_utils enforce)
 endif()
 
-cc_library(kernel_dispatch SRCS kernel_dispatch.cc DEPS pten_tensor device_context kernel_factory)
+cc_library(kernel_dispatch SRCS kernel_dispatch.cc DEPS pten_tensor pten_context kernel_factory)
 
 cc_library(op_meta_info SRCS op_meta_info.cc DEPS pten_tensor)
 
diff --git a/paddle/pten/api/lib/kernel_dispatch.cc b/paddle/pten/api/lib/kernel_dispatch.cc
index 97b3bf281fc8ed45583294644937940373ef9dfc..7930869632a6698581242584efd2b710d6b0b5f5 100644
--- a/paddle/pten/api/lib/kernel_dispatch.cc
+++ b/paddle/pten/api/lib/kernel_dispatch.cc
@@ -51,8 +51,7 @@ std::size_t CountLeadingZeros(uint64_t val) {
 
 }  // namespace detail
 
-paddle::platform::DeviceContext* GetDeviceContextByBackend(
-    pten::Backend backend) {
+pten::DeviceContext* GetDeviceContextByBackend(pten::Backend backend) {
   auto& pool = paddle::platform::DeviceContextPool::Instance();
   return pool.Get(pten::TransToFluidPlace(backend));
 }
diff --git a/paddle/pten/api/lib/kernel_dispatch.h b/paddle/pten/api/lib/kernel_dispatch.h
index 1bba16d107d48789b890d00e10ec942aa416b80d..9c83e0d3d4e9ed794d5d6edce53c317967bf3020 100644
--- a/paddle/pten/api/lib/kernel_dispatch.h
+++ b/paddle/pten/api/lib/kernel_dispatch.h
@@ -21,31 +21,22 @@ limitations under the License. */
 #include "paddle/pten/api/include/tensor.h"
 #include "paddle/pten/api/lib/backend_set.h"
 #include "paddle/pten/api/lib/data_type_set.h"
+#include "paddle/pten/backends/all_context.h"
 #include "paddle/pten/common/data_type.h"
 #include "paddle/pten/common/layout.h"
 
 // TODO(chenweihang): split Key, Kernel, Factory into diff files
 #include "paddle/pten/core/kernel_factory.h"
 
-// See Note [ Why still include the fluid headers? ]
-#include "paddle/fluid/platform/device_context.h"
-
 namespace paddle {
 namespace experimental {
 
-// TODO(shixiaowei): replaced by new DeviceContext later
-using CPUContext = paddle::platform::CPUDeviceContext;
-#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
-using CUDAContext = paddle::platform::CUDADeviceContext;
-#endif
-
 namespace detail {
 BackendSet GetTensorBackendSet(const Tensor& t);
 std::size_t CountLeadingZeros(uint64_t val);
 }  // namespace detail
 
-paddle::platform::DeviceContext* GetDeviceContextByBackend(
-    pten::Backend backend);
+pten::DeviceContext* GetDeviceContextByBackend(pten::Backend backend);
 
 // TODO(chenweihang): support DataLayout and DataType selected
 struct KernelKeySet {
diff --git a/paddle/pten/backends/CMakeLists.txt b/paddle/pten/backends/CMakeLists.txt
index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..af49ad61a48f59fb6559d93985dd27eb3e6ea2fe 100644
--- a/paddle/pten/backends/CMakeLists.txt
+++ b/paddle/pten/backends/CMakeLists.txt
@@ -0,0 +1 @@
+cc_library(pten_context SRCS all_context.cc DEPS device_context)
diff --git a/paddle/pten/backends/all_context.cc b/paddle/pten/backends/all_context.cc
new file mode 100644
index 0000000000000000000000000000000000000000..eb73e41a955081915fdb8d1c10c4e41ccdfeafac
--- /dev/null
+++ b/paddle/pten/backends/all_context.cc
@@ -0,0 +1,17 @@
+/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include "paddle/pten/backends/all_context.h"
+
+namespace pten {}  // namespace pten
diff --git a/paddle/pten/backends/all_context.h b/paddle/pten/backends/all_context.h
new file mode 100644
index 0000000000000000000000000000000000000000..d056af1b3a0d9abb11d9a2a69373c54252213e83
--- /dev/null
+++ b/paddle/pten/backends/all_context.h
@@ -0,0 +1,31 @@
+/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#pragma once
+
+// Note: Some scenarios need to include all types of Context declarations.
+// In order to avoid including the header files of each backend in turn,
+// add this header file
+// Note: Limit the entry of DeviceContext to backends to avoid multiple include
+// path replacement after implementing pten DeviceContext
+
+#include "paddle/pten/backends/cpu/cpu_context.h"
+#include "paddle/pten/backends/cuda/cuda_context.h"
+#include "paddle/pten/backends/npu/npu_context.h"
+#include "paddle/pten/backends/xpu/xpu_context.h"
+
+namespace pten {
+using DeviceContext = paddle::platform::DeviceContext;
+using DeviceContextPool = paddle::platform::DeviceContextPool;
+}  // namespace pten
diff --git a/paddle/pten/backends/cpu/cpu_context.h b/paddle/pten/backends/cpu/cpu_context.h
new file mode 100644
index 0000000000000000000000000000000000000000..b161e98d6361764824ae9273addefef1588d9e26
--- /dev/null
+++ b/paddle/pten/backends/cpu/cpu_context.h
@@ -0,0 +1,22 @@
+/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#pragma once
+
+// See Note [ Why still include the fluid headers? ]
+#include "paddle/fluid/platform/device_context.h"
+
+namespace pten {
+using CPUContext = paddle::platform::CPUDeviceContext;
+}  // namespace pten
diff --git a/paddle/pten/backends/cuda/cuda_context.h b/paddle/pten/backends/cuda/cuda_context.h
new file mode 100644
index 0000000000000000000000000000000000000000..332fdd2fdaf432e8f6363ea577b94ffddd65ca2f
--- /dev/null
+++ b/paddle/pten/backends/cuda/cuda_context.h
@@ -0,0 +1,26 @@
+/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#pragma once
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+
+// See Note [ Why still include the fluid headers? ]
+#include "paddle/fluid/platform/device_context.h"
+
+namespace pten {
+using CUDAContext = paddle::platform::CUDADeviceContext;
+}  // namespace pten
+
+#endif
diff --git a/paddle/pten/backends/npu/npu_context.h b/paddle/pten/backends/npu/npu_context.h
new file mode 100644
index 0000000000000000000000000000000000000000..bb17a1bea6df3c958a199b726f7f7679a6d0fcae
--- /dev/null
+++ b/paddle/pten/backends/npu/npu_context.h
@@ -0,0 +1,26 @@
+/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#pragma once
+
+#ifdef PADDLE_WITH_ASCEND_CL
+
+// See Note [ Why still include the fluid headers? ]
+#include "paddle/fluid/platform/device_context.h"
+
+namespace pten {
+using NPUContext = paddle::platform::NPUDeviceContext;
+}  // namespace pten
+
+#endif  // PADDLE_WITH_ASCEND_CL
diff --git a/paddle/pten/backends/xpu/xpu_context.h b/paddle/pten/backends/xpu/xpu_context.h
new file mode 100644
index 0000000000000000000000000000000000000000..94d2a1532f6365bdb4e916adc54a32f3b5f492f3
--- /dev/null
+++ b/paddle/pten/backends/xpu/xpu_context.h
@@ -0,0 +1,26 @@
+/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#pragma once
+
+#ifdef PADDLE_WITH_XPU
+
+// See Note [ Why still include the fluid headers? ]
+#include "paddle/fluid/platform/device_context.h"
+
+namespace pten {
+using XPUContext = paddle::platform::XPUDeviceContext;
+}  // namespace pten
+
+#endif  // PADDLE_WITH_XPU
diff --git a/paddle/pten/core/CMakeLists.txt b/paddle/pten/core/CMakeLists.txt
index f02678538cb8ea3fd4714745a5e8746c9a1d14a7..63cae6cc70867d247cbeaba71262811ddc073f89 100644
--- a/paddle/pten/core/CMakeLists.txt
+++ b/paddle/pten/core/CMakeLists.txt
@@ -7,7 +7,7 @@ else()
 endif()
 
 cc_library(kernel_factory SRCS kernel_factory.cc DEPS enforce convert_utils)
-cc_library(kernel_context SRCS kernel_context.cc DEPS enforce device_context)
+cc_library(kernel_context SRCS kernel_context.cc DEPS enforce pten_context)
 
 cc_library(tensor_base SRCS tensor_base.cc allocator.cc storage.cc DEPS enforce)
 cc_library(tensor_meta SRCS tensor_meta.cc DEPS enforce)
diff --git a/paddle/pten/core/kernel_utils.h b/paddle/pten/core/kernel_utils.h
index ad7387e15295fd23b215ef7df1dc5191cde7db3d..82ffa573870dfddce1a281a044ace0d08d60050f 100644
--- a/paddle/pten/core/kernel_utils.h
+++ b/paddle/pten/core/kernel_utils.h
@@ -14,6 +14,7 @@
 
 #pragma once
 
+#include "paddle/pten/backends/all_context.h"
 #include "paddle/pten/common/scalar.h"
 #include "paddle/pten/common/scalar_array.h"
 #include "paddle/pten/core/dense_tensor.h"
@@ -21,26 +22,10 @@
 #include "paddle/pten/core/kernel_def.h"
 
 // See Note [ Why still include the fluid headers? ]
-#include "paddle/fluid/platform/device_context.h"
 #include "paddle/fluid/platform/enforce.h"
 
 namespace pten {
 
-// TODO(shixiaowei): replaced by new DeviceContext later
-using CPUContext = paddle::platform::CPUDeviceContext;
-#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
-using CUDAContext = paddle::platform::CUDADeviceContext;
-#endif
-#ifdef PADDLE_WITH_MKLDNN
-using MKLDNNContext = paddle::platform::MKLDNNDeviceContext;
-#endif
-#ifdef PADDLE_WITH_ASCEND_CL
-using NPUContext = paddle::platform::NPUDeviceContext;
-#endif
-#ifdef PADDLE_WITH_XPU
-using XPUContext = paddle::platform::XPUDeviceContext;
-#endif
-
 #define PT_KERNEL(...) \
   ::pten::KernelImpl<decltype(&__VA_ARGS__), &__VA_ARGS__>::Compute