diff --git a/colossalai/kernel/cuda_native/csrc/multihead_attention_1d.cpp b/colossalai/kernel/cuda_native/csrc/multihead_attention_1d.cpp
index b02556f79..166c698f6 100644
--- a/colossalai/kernel/cuda_native/csrc/multihead_attention_1d.cpp
+++ b/colossalai/kernel/cuda_native/csrc/multihead_attention_1d.cpp
@@ -2,8 +2,13 @@
 
 #include <ATen/cuda/CUDAContext.h>
 #include <torch/extension.h>
+#include <torch/torch.h>
 
+#if TORCH_VERSION_MINOR >= 13
+#include <torch/csrc/distributed/c10d/Types.hpp>
+#else
 #include <c10d/Types.hpp>
+#endif
 #include <iostream>
 
 #include "context.h"
diff --git a/colossalai/kernel/cuda_native/csrc/multihead_attention_1d.h b/colossalai/kernel/cuda_native/csrc/multihead_attention_1d.h
index 70b3419d8..db50071b6 100644
--- a/colossalai/kernel/cuda_native/csrc/multihead_attention_1d.h
+++ b/colossalai/kernel/cuda_native/csrc/multihead_attention_1d.h
@@ -4,8 +4,14 @@
 #include <cuda.h>
 #include <cuda_fp16.h>
 #include <cuda_runtime_api.h>
+#include <torch/torch.h>
 
+#if TORCH_VERSION_MINOR >= 13
+#include <torch/csrc/distributed/c10d/ProcessGroup.hpp>
+#else
 #include <c10d/ProcessGroup.hpp>
+#endif
+
 #include <string>
 #include <type_traits>
 
@@ -157,4 +163,4 @@ class MultiHeadAttention {
 
   c10::intrusive_ptr<c10d::ProcessGroup> pg;
   int pg_size;
-};
\ No newline at end of file
+};