diff --git a/xla/service/memory_space_assignment/best_fit_repacker.h b/xla/service/memory_space_assignment/best_fit_repacker.h
index 8fd0f7c1550dc..e22daba33991f 100644
--- a/xla/service/memory_space_assignment/best_fit_repacker.h
+++ b/xla/service/memory_space_assignment/best_fit_repacker.h
@@ -17,6 +17,7 @@ limitations under the License.
 #define XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_BEST_FIT_REPACKER_H_
 
 #include <cstdint>
+#include <utility>
 
 #include "absl/status/statusor.h"
 #include "absl/types/span.h"
diff --git a/xla/service/memory_space_assignment/best_fit_repacker_test.cc b/xla/service/memory_space_assignment/best_fit_repacker_test.cc
index 3003bd69e617e..2b47d1223f800 100644
--- a/xla/service/memory_space_assignment/best_fit_repacker_test.cc
+++ b/xla/service/memory_space_assignment/best_fit_repacker_test.cc
@@ -16,6 +16,8 @@ limitations under the License.
 #include "xla/service/memory_space_assignment/best_fit_repacker.h"
 
 #include <cstdint>
+#include <list>
+#include <vector>
 
 #include "absl/container/flat_hash_map.h"
 #include "absl/types/span.h"
diff --git a/xla/service/memory_space_assignment/slice.h b/xla/service/memory_space_assignment/slice.h
index da3fab681d3f8..f0caa04e92ee4 100644
--- a/xla/service/memory_space_assignment/slice.h
+++ b/xla/service/memory_space_assignment/slice.h
@@ -38,6 +38,7 @@ limitations under the License.
 #define XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_SLICE_H_
 
 #include <cstdint>
+#include <functional>
 #include <ostream>
 #include <string>
 #include <tuple>
diff --git a/xla/service/memory_space_assignment/utils.cc b/xla/service/memory_space_assignment/utils.cc
index b4b37ff0677ba..43f04c263f27e 100644
--- a/xla/service/memory_space_assignment/utils.cc
+++ b/xla/service/memory_space_assignment/utils.cc
@@ -20,6 +20,7 @@ limitations under the License.
 #include <optional>
 
 #include "absl/algorithm/container.h"
+#include "absl/container/flat_hash_map.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
 #include "re2/re2.h"
@@ -30,6 +31,7 @@ limitations under the License.
 #include "xla/hlo/utils/hlo_live_range.h"
 #include "xla/service/heap_simulator/heap_simulator.h"
 #include "xla/service/hlo_value.h"
+#include "xla/service/memory_space_assignment/memory_space_assignment.pb.h"
 #include "xla/shape_util.h"
 #include "xla/util.h"
 #include "tsl/platform/statusor.h"