Add buffer allocation telemetry to TPU Raiden KV cache manager.

tpu_sync/kv_cache/BUILD

@@ -240,6 +240,7 @@ cc_library(
         "//tpu_sync/transport:block_transport",
         "//tpu_sync/transport:block_transport_delegate",
         "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/cleanup",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/functional:any_invocable",
         "@com_google_absl//absl/log",

tpu_sync/kv_cache/kv_cache_manager_base.cc

@@ -34,6 +34,7 @@
 #include <vector>
 
 #include "absl/base/thread_annotations.h"
+#include "absl/cleanup/cleanup.h"
 #include "absl/log/log.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
@@ -272,6 +273,7 @@ KVCacheManagerBase::KVCacheManagerBase(
 
   layers_.reserve(num_layers_);
   buffer_holds_.reserve(num_layers_);
+  size_t total_host_dram_bytes = 0;
 
   for (size_t layer_idx = 0; layer_idx < num_layers_; ++layer_idx) {
     const auto& dst_buffers = layer_buffers[layer_idx];
@@ -349,18 +351,24 @@ KVCacheManagerBase::KVCacheManagerBase(
                 << shard_info.host_size;
       }
 
+      total_host_dram_bytes += shard_info.host_size;
       device_info.holds.push_back(dst_buffer);
       layer_info.shards.push_back(std::move(shard_info));
     }
     layers_.push_back(std::move(layer_info));
     buffer_holds_.push_back(std::move(device_info));
   }
+  {
+    absl::MutexLock lock(allocated_host_dram_bytes_mu_);
+    allocated_host_dram_bytes_ = total_host_dram_bytes;
+  }
 
   constexpr size_t kPoolSize = 4;
   dma_pool_ = std::make_unique<NumaThreadPool>(kPoolSize);
   push_pool_ = std::make_shared<NumaThreadPool>(kPoolSize);
   pull_pool_ = std::make_unique<NumaThreadPool>(kPoolSize);
   InitBackgroundWorker();
+  UpdateAllocatedOccupancyMetric();
 }
 
 KVCacheManagerBase::KVCacheManagerBase(
@@ -376,6 +384,7 @@ KVCacheManagerBase::KVCacheManagerBase(
   semaphore_ = std::make_unique<xla::Semaphore>(std::max<int>(4, parallelism));
 
   layers_.reserve(num_layers_);
+  size_t total_host_dram_bytes = 0;
   for (size_t layer_idx = 0; layer_idx < num_layers_; ++layer_idx) {
     LayerInfoBase layer_info;
     layer_info.shards.reserve(num_shards_);
@@ -422,16 +431,22 @@ KVCacheManagerBase::KVCacheManagerBase(
         shard_info.host_size = alloc_size;
       }
 
+      total_host_dram_bytes += shard_info.host_size;
       layer_info.shards.push_back(std::move(shard_info));
     }
     layers_.push_back(std::move(layer_info));
   }
+  {
+    absl::MutexLock lock(allocated_host_dram_bytes_mu_);
+    allocated_host_dram_bytes_ = total_host_dram_bytes;
+  }
   constexpr size_t kPoolSize = 4;
   dma_pool_ = std::make_unique<NumaThreadPool>(kPoolSize);
   push_pool_ = std::make_shared<NumaThreadPool>(kPoolSize);
   pull_pool_ = std::make_unique<NumaThreadPool>(kPoolSize);
   InitTransportServer();
   InitBackgroundWorker();
+  UpdateAllocatedOccupancyMetric();
 }
 
 void KVCacheManagerBase::InitBackgroundWorker() {
@@ -1195,25 +1210,35 @@ absl::StatusOr<raiden::PjRtCopyFuture> KVCacheManagerBase::H2hReadExplicit(
 void KVCacheManagerBase::SetExternalHostBuffer(
     const std::vector<raiden::BufferHoldAndAlias>& buffer_holds) {
   size_t idx = 0;
+  size_t old_total = 0;
+  size_t new_total = 0;
   for (size_t l = 0; l < num_layers_; ++l) {
     for (size_t sh = 0; sh < num_shards_; ++sh) {
       if (idx < buffer_holds.size()) {
         void* host_ptr = buffer_holds[idx].GetHostPointer();
         if (host_ptr) {
+          old_total += layers_[l].shards[sh].host_size;
+          const size_t new_size = buffer_holds[idx].GetOnDeviceSizeInBytes();
+          new_total += new_size;
           layers_[l].shards[sh].host_ptr = reinterpret_cast<uint8_t*>(host_ptr);
-          layers_[l].shards[sh].host_size =
-              buffer_holds[idx].GetOnDeviceSizeInBytes();
+          layers_[l].shards[sh].host_size = new_size;
         }
         idx++;
       }
     }
   }
+  {
+    absl::MutexLock lock(allocated_host_dram_bytes_mu_);
+    allocated_host_dram_bytes_ =
+        allocated_host_dram_bytes_ - old_total + new_total;
+  }
   // Host geometry changed: drop any lazily built implicit pools so the next
   // pool access rebuilds them against the new buffers.
   absl::MutexLock l(pools_mu_);
   if (!explicit_pools_) {
     pools_.clear();
   }
+  UpdateAllocatedOccupancyMetric();
 }
 
 absl::Status KVCacheManagerBase::H2dDirect(
@@ -1569,11 +1594,22 @@ absl::Status KVCacheManagerBase::EnsureHostMirrorCovers(size_t storage_idx,
   if (storage_idx >= layers_.size() || needed_bytes <= 0) {
     return absl::OkStatus();
   }
+  size_t old_total = 0;
+  size_t new_total = 0;
+  auto update_bytes_on_exit = absl::MakeCleanup([this, &old_total, &new_total] {
+    if (old_total != new_total) {
+      absl::MutexLock lock(allocated_host_dram_bytes_mu_);
+      allocated_host_dram_bytes_ =
+          allocated_host_dram_bytes_ - old_total + new_total;
+    }
+  });
+
   for (auto& shard_info : layers_[storage_idx].shards) {
     if (static_cast<int64_t>(shard_info.host_size) >= needed_bytes) {
       continue;
     }
     const size_t alloc_size = static_cast<size_t>(needed_bytes);
+    const size_t prev_size = shard_info.host_size;
     if (host_allocator_) {
       ASSIGN_OR_RETURN(HostBufferAllocation allocation,
                        host_allocator_(alloc_size, nullptr));
@@ -1589,6 +1625,8 @@ absl::Status KVCacheManagerBase::EnsureHostMirrorCovers(size_t storage_idx,
       shard_info.host_size = allocation.size;
       shard_info.host_owner = std::move(allocation.owner);
       shard_info.owned_host_buffer = {nullptr, [](void*) {}};
+      old_total += prev_size;
+      new_total += allocation.size;
     } else {
       void* ptr = nullptr;
       if (posix_memalign(&ptr, 64, alloc_size) != 0) {
@@ -1605,6 +1643,8 @@ absl::Status KVCacheManagerBase::EnsureHostMirrorCovers(size_t storage_idx,
       shard_info.host_ptr = shard_info.owned_host_buffer.get();
       shard_info.host_size = alloc_size;
       shard_info.host_owner.reset();
+      old_total += prev_size;
+      new_total += alloc_size;
     }
   }
   return absl::OkStatus();
@@ -1662,9 +1702,12 @@ absl::Status KVCacheManagerBase::RegisterPools(std::vector<PoolSpec> pools) {
       }
     }
   }
-  absl::MutexLock l(pools_mu_);
-  pools_ = std::move(pools);
-  explicit_pools_ = true;
+  {
+    absl::MutexLock l(pools_mu_);
+    pools_ = std::move(pools);
+    explicit_pools_ = true;
+  }
+  UpdateAllocatedOccupancyMetric();
   return absl::OkStatus();
 }
 
@@ -2567,5 +2610,20 @@ absl::StatusOr<raiden::PjRtCopyFuture> KVCacheManagerBase::D2h(
                          target_shard_idx);
 }
 
+size_t KVCacheManagerBase::GetAllocatedHostDramBytes() const {
+  absl::MutexLock lock(allocated_host_dram_bytes_mu_);
+  return allocated_host_dram_bytes_;
+}
+
+void KVCacheManagerBase::UpdateAllocatedOccupancyMetric() const {
+  if (!telemetry::RaidenMetricStore::GetGlobalMetricStore().HasBackends()) {
+    return;
+  }
+  const size_t total_host_dram = GetAllocatedHostDramBytes();
+  telemetry::RaidenMetricStore::GetGlobalMetricStore().SetGauge(
+      telemetry::metric_names::kBufferAllocatedBytes, {},
+      static_cast<double>(total_host_dram));
+}
+
 }  // namespace kv_cache
 }  // namespace tpu_raiden

tpu_sync/kv_cache/kv_cache_manager_base.h

@@ -15,6 +15,7 @@
 #ifndef THIRD_PARTY_TPU_RAIDEN_KV_CACHE_KV_CACHE_MANAGER_BASE_H_
 #define THIRD_PARTY_TPU_RAIDEN_KV_CACHE_KV_CACHE_MANAGER_BASE_H_
 
+#include <atomic>
 #include <cstddef>
 #include <cstdint>
 #include <functional>
@@ -391,6 +392,9 @@ class KVCacheManagerBase : public tpu_raiden::RaidenManagerBase {
   uint8_t* GetBlockHostPointer(size_t layer_idx, size_t shard_idx,
                                int block_id) override;
 
+  // Returns the total number of bytes allocated in host DRAM.
+  size_t GetAllocatedHostDramBytes() const;
+
  protected:
   const PJRT_Api* c_api_ = nullptr;
   const PJRT_RawBuffer_Extension* extension_ = nullptr;
@@ -485,6 +489,9 @@ class KVCacheManagerBase : public tpu_raiden::RaidenManagerBase {
       std::optional<size_t> layer_idx = std::nullopt,
       std::optional<size_t> shard_idx = std::nullopt);
 
+  // Updates the allocated buffer occupancy metrics.
+  void UpdateAllocatedOccupancyMetric() const;
+
  private:
   // Asynchronous on-chip H2D offload enqueued to the background worker queue.
   virtual absl::StatusOr<raiden::PjRtCopyFuture> H2dAsyncDispatch(
@@ -575,6 +582,13 @@ class KVCacheManagerBase : public tpu_raiden::RaidenManagerBase {
   // RAIDEN_ENABLE_ASYNC_DISPATCH environment variable.
   bool enable_background_ = false;
 
+  // Mutex guarding allocated_host_dram_bytes_.
+  mutable absl::Mutex allocated_host_dram_bytes_mu_;
+
+  // Running total of allocated host DRAM bytes.
+  size_t allocated_host_dram_bytes_
+      ABSL_GUARDED_BY(allocated_host_dram_bytes_mu_) = 0;
+
   // Initializes background worker thread if RAIDEN_ENABLE_ASYNC_DISPATCH is
   // enabled.
   void InitBackgroundWorker();

tpu_sync/kv_cache/kv_cache_manager_test.cc

@@ -31,10 +31,12 @@
 #include "absl/synchronization/mutex.h"
 #include "absl/strings/string_view.h"
 #include "tpu_sync/core/raw_transfer_core.h"
+#include "tpu_sync/core/raiden_manager_base.h"
 #include "tpu_sync/kv_cache/kv_cache_manager_base.h"
 #include "tpu_sync/rpc/raiden_service.pb.h"
 #include "tpu_sync/telemetry/metrics_api.h"
 #include "tpu_sync/telemetry/metrics_backend.h"
+#include "tpu_sync/telemetry/mock_metrics_backend.h"
 #include "tpu_sync/transport/block_transport.h"
 
 namespace tpu_raiden {
@@ -65,6 +67,9 @@ class TestKVCacheManager : public KVCacheManagerBase {
     buffer_holds_[layer_idx].physical_size = physical_size;
     major_dim_size_ = major_dim_size;
   }
+
+  using KVCacheManagerBase::layers_;
+  using KVCacheManagerBase::UpdateAllocatedOccupancyMetric;
 };
 
 // Pool with one strided live region per block: live [0, 32) and [64, 96)
@@ -1361,6 +1366,71 @@ TEST(KVCacheManagerTest, D2hWritePipelinedTelemetryBatchObservation) {
   EXPECT_OK(res.Await());
 }
 
+TEST(KVCacheManagerTest, BufferAllocatedHostDramTelemetry) {
+  auto mock_backend = std::make_unique<telemetry::MockMetricsBackend>();
+  telemetry::MockMetricsBackend* raw_mock = mock_backend.get();
+  telemetry::ScopedMetricsBackendReset scoped_reset(std::move(mock_backend));
+
+  // Expect initial gauge sets in constructor:
+  // num_layers = 2, num_shards = 2, slice_byte_size = 128, host_blocks = 4
+  // total allocated host dram = 2 * 2 * (4 * 128) = 2048 bytes
+  EXPECT_CALL(
+      *raw_mock,
+      SetGauge(testing::Eq(telemetry::metric_names::kBufferAllocatedBytes),
+               testing::IsEmpty(), testing::DoubleEq(2048.0)))
+      .Times(testing::AtLeast(1));
+
+  TestKVCacheManager manager(/*num_layers=*/2, /*num_shards=*/2,
+                             /*slice_byte_size=*/128,
+                             /*host_blocks=*/4);
+
+  EXPECT_EQ(manager.GetAllocatedHostDramBytes(), 2048);
+}
+
+TEST(KVCacheManagerTest, BufferAllocatedHostDramMaintenance) {
+  auto mock_backend = std::make_unique<telemetry::MockMetricsBackend>();
+  telemetry::MockMetricsBackend* raw_mock = mock_backend.get();
+  telemetry::ScopedMetricsBackendReset scoped_reset(std::move(mock_backend));
+
+  // Initial gauge set in constructor (2 * 2 * 4 * 128 = 2048 bytes).
+  EXPECT_CALL(
+      *raw_mock,
+      SetGauge(testing::Eq(telemetry::metric_names::kBufferAllocatedBytes),
+               testing::IsEmpty(), testing::DoubleEq(2048.0)))
+      .Times(1);
+
+  TestKVCacheManager manager(/*num_layers=*/2, /*num_shards=*/2,
+                             /*slice_byte_size=*/128,
+                             /*host_blocks=*/4);
+
+  EXPECT_EQ(manager.GetAllocatedHostDramBytes(), 2048);
+
+  // Calling RegisterPools grows layer 0 host mirror from 4*128 to 8*128 bytes
+  // per shard (2 shards = +1024 bytes -> 3072 bytes total).
+  EXPECT_CALL(
+      *raw_mock,
+      SetGauge(testing::Eq(telemetry::metric_names::kBufferAllocatedBytes),
+               testing::IsEmpty(), testing::DoubleEq(3072.0)))
+      .Times(1);
+  manager.SetLayerPhysicalSizeForTest(/*layer_idx=*/0,
+                                      /*physical_size=*/1024,
+                                      /*major_dim_size=*/1);
+  absl::Status status = manager.RegisterPools(
+      {DensePool("kind_a", /*storage_index=*/0, /*base_offset=*/0,
+                 /*stride=*/128, /*num_blocks=*/8)});
+  ASSERT_TRUE(status.ok()) << status.ToString();
+  EXPECT_EQ(manager.GetAllocatedHostDramBytes(), 3072);
+
+  // Subsequent call to UpdateAllocatedOccupancyMetric maintains and reports
+  // the running value (3072).
+  EXPECT_CALL(
+      *raw_mock,
+      SetGauge(testing::Eq(telemetry::metric_names::kBufferAllocatedBytes),
+               testing::IsEmpty(), testing::DoubleEq(3072.0)))
+      .Times(1);
+  manager.UpdateAllocatedOccupancyMetric();
+}
+
 }  // namespace
 }  // namespace kv_cache
 }  // namespace tpu_raiden

tpu_sync/telemetry/metrics_api_test.cc

@@ -165,6 +165,20 @@ TEST_F(MetricsApiTest, MetricMetadataConstants) {
       "milliseconds, including setup delays.");
   EXPECT_EQ(metric_metadata::kTransferDurationMs.type, MetricType::kHistogram);
 
+  // BufferAllocatedBytes
+  EXPECT_EQ(metric_names::kBufferAllocatedBytes, "buffer_allocated_bytes");
+  EXPECT_EQ(metric_descriptions::kBufferAllocatedBytes,
+            "Current host DRAM buffer capacity allocated in bytes for KV cache "
+            "staging "
+            "across all layers and shards.");
+  EXPECT_EQ(metric_metadata::kBufferAllocatedBytes.name,
+            "buffer_allocated_bytes");
+  EXPECT_EQ(metric_metadata::kBufferAllocatedBytes.description,
+            "Current host DRAM buffer capacity allocated in bytes for KV cache "
+            "staging "
+            "across all layers and shards.");
+  EXPECT_EQ(metric_metadata::kBufferAllocatedBytes.type, MetricType::kGauge);
+
   // Direction Labels
   EXPECT_EQ(metric_labels::kDirection, "direction");
   EXPECT_EQ(metric_labels::kDirectionPush, "push");
@@ -181,7 +195,8 @@ TEST_F(MetricsApiTest, MetricMetadataConstants) {
                           metric_metadata::kTransferFailuresTotal,
                           metric_metadata::kTransferDurationMs,
                           metric_metadata::kH2dTransferTimeMs,
-                          metric_metadata::kD2hTransferTimeMs));
+                          metric_metadata::kD2hTransferTimeMs,
+                          metric_metadata::kBufferAllocatedBytes));
 }
 
 TEST_F(MetricsApiTest, FastPathExitWhenNoBackends) {
@@ -317,7 +332,8 @@ TEST_F(MetricsApiTest, GetMetricMetadataReturnsAllMetricsWhenBackendsActive) {
                                   metric_metadata::kTransferFailuresTotal,
                                   metric_metadata::kTransferDurationMs,
                                   metric_metadata::kH2dTransferTimeMs,
-                                  metric_metadata::kD2hTransferTimeMs));
+                                  metric_metadata::kD2hTransferTimeMs,
+                                  metric_metadata::kBufferAllocatedBytes));
 }
 
 TEST_F(MetricsApiTest, GetMetricMetadataEmptyWhenNoBackends) {