From 581c4cb1439e1bc982bd87ab38feff20b02d6fff Mon Sep 17 00:00:00 2001
From: Danna Wang <dannawang@google.com>
Date: Mon, 24 Aug 2026 14:58:01 -0700
Subject: [PATCH] Fix compilation errors in KVCacheStoreMock.

PiperOrigin-RevId: 970100062
---
 tpu_sync/frameworks/torch/BUILD | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tpu_sync/frameworks/torch/BUILD b/tpu_sync/frameworks/torch/BUILD
index 5f21dab4..5f0e7dfd 100644
--- a/tpu_sync/frameworks/torch/BUILD
+++ b/tpu_sync/frameworks/torch/BUILD
@@ -429,6 +429,7 @@ cc_library(
     tags = ["nobuilder"],
     visibility = ["//visibility:public"],
     deps = [
+        ":torch_tpu_device_buffer_headers",
         ":torch_tpu_tensor_to_buffer_headers",
         ":torch_utils_mock",
         "//tpu_sync/weight_sync:weight_synchronizer_base",