diff --git a/xla/stream_executor/cuda/cuda_kernel.cc b/xla/stream_executor/cuda/cuda_kernel.cc
index 265c885c69e94..48697624cbb60 100644
--- a/xla/stream_executor/cuda/cuda_kernel.cc
+++ b/xla/stream_executor/cuda/cuda_kernel.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include "xla/stream_executor/cuda/cuda_kernel.h"
 
+#include <atomic>
 #include <cstddef>
 #include <cstdint>
 #include <memory>
@@ -26,6 +27,7 @@ limitations under the License.
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
+#include "absl/synchronization/mutex.h"
 #include "third_party/gpus/cuda/include/cuda.h"
 #include "xla/stream_executor/activate_context.h"
 #include "xla/stream_executor/cuda/cuda_status.h"