diff --git a/xla/backends/cpu/runtime/BUILD b/xla/backends/cpu/runtime/BUILD
index da32c5e807d99..e316e3813062b 100644
--- a/xla/backends/cpu/runtime/BUILD
+++ b/xla/backends/cpu/runtime/BUILD
@@ -868,6 +868,7 @@ cc_library(
         ":thunk",
         "//xla:shape_util",
         "//xla:status_macros",
+        "//xla:xla_data_proto_cc",
         "//xla/runtime:buffer_use",
         "//xla/runtime:device_id",
         "//xla/service:buffer_assignment",
diff --git a/xla/backends/cpu/runtime/logical_id_thunk.cc b/xla/backends/cpu/runtime/logical_id_thunk.cc
index ef33770ab5cd4..74a295ab56dda 100644
--- a/xla/backends/cpu/runtime/logical_id_thunk.cc
+++ b/xla/backends/cpu/runtime/logical_id_thunk.cc
@@ -33,6 +33,7 @@ limitations under the License.
 #include "xla/stream_executor/device_address.h"
 #include "xla/tsl/concurrency/async_value_ref.h"
 #include "xla/tsl/platform/statusor.h"
+#include "xla/xla_data.pb.h"
 
 namespace xla::cpu::internal {