diff --git a/tpu_sync/kv_cache/BUILD b/tpu_sync/kv_cache/BUILD index e0650771..2bc32e8c 100644 --- a/tpu_sync/kv_cache/BUILD +++ b/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", diff --git a/tpu_sync/kv_cache/kv_cache_manager_base.cc b/tpu_sync/kv_cache/kv_cache_manager_base.cc index cae1e4c8..098bcf4b 100644 --- a/tpu_sync/kv_cache/kv_cache_manager_base.cc +++ b/tpu_sync/kv_cache/kv_cache_manager_base.cc @@ -34,6 +34,7 @@ #include #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(kPoolSize); push_pool_ = std::make_shared(kPoolSize); pull_pool_ = std::make_unique(kPoolSize); InitBackgroundWorker(); + UpdateAllocatedOccupancyMetric(); } KVCacheManagerBase::KVCacheManagerBase( @@ -376,6 +384,7 @@ KVCacheManagerBase::KVCacheManagerBase( semaphore_ = std::make_unique(std::max(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(kPoolSize); push_pool_ = std::make_shared(kPoolSize); pull_pool_ = std::make_unique(kPoolSize); InitTransportServer(); InitBackgroundWorker(); + UpdateAllocatedOccupancyMetric(); } void KVCacheManagerBase::InitBackgroundWorker() { @@ -1195,25 +1210,35 @@ absl::StatusOr KVCacheManagerBase::H2hReadExplicit( void KVCacheManagerBase::SetExternalHostBuffer( const std::vector& 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(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(shard_info.host_size) >= needed_bytes) { continue; } const size_t alloc_size = static_cast(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 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 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(total_host_dram)); +} + } // namespace kv_cache } // namespace tpu_raiden diff --git a/tpu_sync/kv_cache/kv_cache_manager_base.h b/tpu_sync/kv_cache/kv_cache_manager_base.h index 858c3071..8758624b 100644 --- a/tpu_sync/kv_cache/kv_cache_manager_base.h +++ b/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 #include #include #include @@ -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 layer_idx = std::nullopt, std::optional 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 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(); diff --git a/tpu_sync/kv_cache/kv_cache_manager_test.cc b/tpu_sync/kv_cache/kv_cache_manager_test.cc index eaa9128e..76cd6cf7 100644 --- a/tpu_sync/kv_cache/kv_cache_manager_test.cc +++ b/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* 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* 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 diff --git a/tpu_sync/telemetry/metrics_api_test.cc b/tpu_sync/telemetry/metrics_api_test.cc index 989c09bd..1e253d73 100644 --- a/tpu_sync/telemetry/metrics_api_test.cc +++ b/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) { diff --git a/tpu_sync/telemetry/metrics_backend.h b/tpu_sync/telemetry/metrics_backend.h index 6ffe1d28..afb0f799 100644 --- a/tpu_sync/telemetry/metrics_backend.h +++ b/tpu_sync/telemetry/metrics_backend.h @@ -61,6 +61,8 @@ inline constexpr absl::string_view kTransferDurationMs = "transfer_duration_ms"; inline constexpr absl::string_view kH2dTransferTimeMs = "h2d_transfer_time_ms"; inline constexpr absl::string_view kD2hTransferTimeMs = "d2h_transfer_time_ms"; +inline constexpr absl::string_view kBufferAllocatedBytes = + "buffer_allocated_bytes"; } // namespace metric_names @@ -75,6 +77,9 @@ inline constexpr absl::string_view kTransferFailuresTotal = inline constexpr absl::string_view kTransferDurationMs = "Measures End-to-End (E2E) latency bound around the entire request in " "milliseconds, including setup delays."; +inline constexpr absl::string_view kBufferAllocatedBytes = + "Current host DRAM buffer capacity allocated in bytes for KV cache staging " + "across all layers and shards."; inline constexpr absl::string_view kH2dTransferTimeMs = "Host-to-Device transfer latency in milliseconds."; @@ -115,9 +120,15 @@ inline constexpr MetricMetadata kD2hTransferTimeMs{ .description = metric_descriptions::kD2hTransferTimeMs, .type = MetricType::kHistogram}; +inline constexpr MetricMetadata kBufferAllocatedBytes{ + .name = metric_names::kBufferAllocatedBytes, + .description = metric_descriptions::kBufferAllocatedBytes, + .type = MetricType::kGauge}; + inline constexpr MetricMetadata kAllMetrics[] = { - kSentBytesTotal, kReceivedBytesTotal, kTransferFailuresTotal, - kTransferDurationMs, kH2dTransferTimeMs, kD2hTransferTimeMs, + kSentBytesTotal, kReceivedBytesTotal, kTransferFailuresTotal, + kTransferDurationMs, kH2dTransferTimeMs, kD2hTransferTimeMs, + kBufferAllocatedBytes, }; } // namespace metric_metadata