[TPU/Raiden/StagingHostRam] Add num_staging_layers constructor parameter to PyTorch KVCacheManager, replacing TPU_RAIDEN_NUM_STAGING_LAYERS

tpu_sync/api/torch/kv_cache_manager.py

@@ -118,6 +118,7 @@ def __init__(
       worker_id: Optional[str] = None,
       enable_shm: bool = False,
       backend_configs: Optional[Sequence[Any]] = None,
+      num_staging_layers: Optional[int] = None,
   ):
     """Instantiates the TransferEngine-based KVCacheManager.
 
@@ -157,9 +158,25 @@ def __init__(
         and set `<axis>_rank` to this worker's rank in [0, size). Duplicate or
         missing ranks cause wrong or missed recalls. An invalid config aborts
         the process (LOG(FATAL)) at registration.
+      num_staging_layers: Optional number of host staging layer buffers to
+        allocate, in [1, num_layers]. Layer-major transfers lease these buffers
+        per layer, bounding host DRAM to num_staging_layers x num_slots x
+        max_blocks blocks instead of a full per-layer mirror. None keeps one
+        host buffer per layer. Only valid on the (max_blocks, num_slots) path.
     """
     self._admission_summary = None
     impl = _torch_impl()
+    if num_staging_layers is not None:
+      if host_blocks_to_allocate is not None:
+        raise ValueError(
+            "num_staging_layers is not supported with host_blocks_to_allocate."
+        )
+      # Checked here, before the C++ constructor pins the TPU buffers.
+      if not 1 <= num_staging_layers <= len(kv_caches):
+        raise ValueError(
+            f"num_staging_layers={num_staging_layers} must be in"
+            f" [1, {len(kv_caches)}] (number of kv_caches layers)."
+        )
     if host_blocks_to_allocate is not None:
       self._impl = impl.KVCacheManager(
           kv_caches,
@@ -195,6 +212,7 @@ def __init__(
           raiden_controller_address=raiden_controller_address,
           worker_id=worker_id,
           enable_shm=enable_shm,
+          num_staging_layers=num_staging_layers,
       )
     if backend_configs:
       self._impl.register_kv_backends(list(backend_configs))

tpu_sync/frameworks/torch/kv_cache_manager.cc

@@ -119,7 +119,8 @@ TorchKVCacheManager::TorchKVCacheManager(
     std::optional<int> host_blocks_to_allocate, bool unsafe_skip_buffer_lock,
     int parallelism, int64_t node_id, int64_t local_control_port,
     int64_t max_blocks, int64_t num_slots, double timeout_s,
-    std::vector<at::Tensor> kv_caches, bool enable_shm)
+    std::vector<at::Tensor> kv_caches, bool enable_shm,
+    std::optional<int> num_staging_layers)
     : KVCacheManagerWithTransfer(
           unpacked.buffers,
           unpacked.has_logical_metadata ? unpacked.logical_slice_byte_size : 0,
@@ -141,7 +142,8 @@ TorchKVCacheManager::TorchKVCacheManager(
                   : unpacked.buffers[0][0]
                         .buffer->GetOnDeviceSizeInBytes()
                         .value_or(0)),
-          node_id, local_control_port, max_blocks, num_slots, timeout_s),
+          node_id, local_control_port, max_blocks, num_slots, timeout_s,
+          /*metrics_collector=*/nullptr, num_staging_layers),
       kv_caches_(std::move(kv_caches)),
       // Move the keep-alive refs in AFTER the base ctor has acquired the
       // buffers; they pin the materialized device buffers for our lifetime.
@@ -153,14 +155,15 @@ TorchKVCacheManager::Create(const std::vector<at::Tensor>& kv_caches,
                             int64_t max_blocks, int64_t num_slots,
                             double timeout_s, bool unsafe_skip_buffer_lock,
                             int parallelism, std::optional<int> listener_port,
-                            bool enable_shm) {
+                            bool enable_shm,
+                            std::optional<int> num_staging_layers) {
   try {
     auto mgr = absl::WrapUnique(new TorchKVCacheManager(
         UnpackLayers(SingleShardLayers(kv_caches)),
         /*local_port=*/std::nullopt,
         /*host_blocks_to_allocate=*/std::nullopt, unsafe_skip_buffer_lock,
         parallelism, node_id, local_control_port, max_blocks, num_slots,
-        timeout_s, kv_caches, enable_shm));
+        timeout_s, kv_caches, enable_shm, num_staging_layers));
     if (listener_port) {
       mgr->listener_ = std::make_unique<kv_cache::KVCacheListener>(
           mgr.get(), *listener_port);
@@ -314,12 +317,14 @@ KVCacheManager::CreateWithNodeId(
     double timeout_s, bool unsafe_skip_buffer_lock, int parallelism,
     std::optional<int> listener_port, int raiden_worker_port,
     std::optional<std::string> raiden_controller_address,
-    std::optional<std::string> worker_id, bool enable_shm) {
-  ABSL_ASSIGN_OR_RETURN(std::unique_ptr<TorchKVCacheManager> torch_manager,
-                        TorchKVCacheManager::Create(
-                            kv_caches, node_id, local_control_port, max_blocks,
-                            num_slots, timeout_s, unsafe_skip_buffer_lock,
-                            parallelism, listener_port, enable_shm));
+    std::optional<std::string> worker_id, bool enable_shm,
+    std::optional<int> num_staging_layers) {
+  ABSL_ASSIGN_OR_RETURN(
+      std::unique_ptr<TorchKVCacheManager> torch_manager,
+      TorchKVCacheManager::Create(
+          kv_caches, node_id, local_control_port, max_blocks, num_slots,
+          timeout_s, unsafe_skip_buffer_lock, parallelism, listener_port,
+          enable_shm, num_staging_layers));
   auto manager = absl::WrapUnique(new KVCacheManager(std::move(torch_manager)));
   ABSL_RETURN_IF_ERROR(manager->StartGrpcServer(
       raiden_worker_port, std::move(raiden_controller_address),

tpu_sync/frameworks/torch/kv_cache_manager.h

@@ -60,7 +60,8 @@ class TorchKVCacheManager : public KVCacheManagerWithTransfer {
       const std::vector<at::Tensor>& kv_caches, int64_t node_id,
       int64_t local_control_port, int64_t max_blocks, int64_t num_slots,
       double timeout_s, bool unsafe_skip_buffer_lock, int parallelism = 4,
-      std::optional<int> listener_port = std::nullopt, bool enable_shm = false);
+      std::optional<int> listener_port = std::nullopt, bool enable_shm = false,
+      std::optional<int> num_staging_layers = std::nullopt);
 
   ~TorchKVCacheManager() override;
 
@@ -129,7 +130,8 @@ class TorchKVCacheManager : public KVCacheManagerWithTransfer {
                       bool unsafe_skip_buffer_lock, int parallelism,
                       int64_t node_id, int64_t local_control_port,
                       int64_t max_blocks, int64_t num_slots, double timeout_s,
-                      std::vector<at::Tensor> kv_caches, bool enable_shm);
+                      std::vector<at::Tensor> kv_caches, bool enable_shm,
+                      std::optional<int> num_staging_layers = std::nullopt);
 
   std::vector<at::Tensor> kv_caches_;
   // Keep-alives for the materialized device buffers backing the manager.
@@ -157,8 +159,10 @@ class KVCacheManager {
   // Creates a transfer-enabled manager over the flat list of per-layer
   // tensors in |kv_caches| (single shard per layer). |enable_shm| opts this
   // manager's host buffers into the shared-memory segments named by
-  // RAIDEN_SHM_KEY; without it the env var is ignored. Failure modes as for
-  // the cache-only Create().
+  // RAIDEN_SHM_KEY; without it the env var is ignored. |num_staging_layers|
+  // bounds the host mirror to that many host layer buffers (see
+  // KVCacheManagerBase). Failure modes as for the cache-only Create(), plus
+  // InvalidArgument for an out-of-range |num_staging_layers|.
   static absl::StatusOr<std::unique_ptr<KVCacheManager>> CreateWithNodeId(
       const std::vector<at::Tensor>& kv_caches, int64_t node_id,
       int64_t local_control_port, int64_t max_blocks, int64_t num_slots,
@@ -167,7 +171,8 @@ class KVCacheManager {
       int raiden_worker_port = 0,
       std::optional<std::string> raiden_controller_address = std::nullopt,
       std::optional<std::string> worker_id = std::nullopt,
-      bool enable_shm = false);
+      bool enable_shm = false,
+      std::optional<int> num_staging_layers = std::nullopt);
 
   ~KVCacheManager();
 

tpu_sync/frameworks/torch/tpu_raiden_torch_module.cc

@@ -270,7 +270,8 @@ NB_MODULE(_tpu_raiden_torch, m) {
            nb::arg("parallelism") = 4, nb::arg("listener_port") = nb::none(),
            nb::arg("raiden_worker_port") = 0,
            nb::arg("raiden_controller_address") = nb::none(),
-           nb::arg("worker_id") = nb::none(), nb::arg("enable_shm") = false)
+           nb::arg("worker_id") = nb::none(), nb::arg("enable_shm") = false,
+           nb::arg("num_staging_layers") = nb::none())
       .def("node_id", &KVCacheManager::node_id)
       .def("get_raiden_worker_port", &KVCacheManager::GetRaidenWorkerPort)
       .def(

tpu_sync/kv_cache/kv_cache_manager_base.cc

@@ -45,7 +45,6 @@
 #include "absl/status/statusor.h"
 #include "absl/strings/ascii.h"
 #include "absl/strings/match.h"
-#include "absl/strings/numbers.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_join.h"
 #include "absl/strings/string_view.h"
@@ -415,17 +414,24 @@ std::vector<transport::ShardInfo> BuildDefaultShards(size_t num_shards) {
   return shards;
 }
 
-// TPU_RAIDEN_NUM_STAGING_LAYERS: number of host layer buffers to allocate.
-// Unset, unparsable, zero, or >= |num_layers| keeps the full per-layer mirror.
-size_t ResolveEnvNumStagingLayers(size_t num_layers) {
-  const char* raw_env = std::getenv("TPU_RAIDEN_NUM_STAGING_LAYERS");
-  size_t requested = 0;
-  if (raw_env == nullptr || !absl::SimpleAtoi(raw_env, &requested) ||
-      requested == 0 || requested >= num_layers) {
+// Number of host layer buffers to allocate: |num_staging_layers| when set,
+// which must lie in [1, |num_layers|]; otherwise the full per-layer mirror.
+size_t ResolveNumStagingLayers(size_t num_layers,
+                               std::optional<int> num_staging_layers) {
+  if (!num_staging_layers.has_value()) {
     return num_layers;
   }
-  LOG(INFO) << "KVCacheManagerBase: allocating " << requested << "/"
-            << num_layers << " host staging layer buffers";
+  if (*num_staging_layers < 1 ||
+      static_cast<size_t>(*num_staging_layers) > num_layers) {
+    throw std::invalid_argument(
+        absl::StrCat("num_staging_layers (", *num_staging_layers,
+                     ") must be in [1, ", num_layers, "]"));
+  }
+  const size_t requested = static_cast<size_t>(*num_staging_layers);
+  if (requested < num_layers) {
+    LOG(INFO) << "KVCacheManagerBase: allocating " << requested << "/"
+              << num_layers << " host staging layer buffers";
+  }
   return requested;
 }
 
@@ -439,7 +445,7 @@ absl::Status CheckAllLayersHaveHostBuffers(
     return absl::FailedPreconditionError(absl::StrCat(
         direction, " of every layer needs a host buffer per layer, but only ",
         num_host_layers, " of ", num_layers,
-        " are allocated (TPU_RAIDEN_NUM_STAGING_LAYERS)"));
+        " are allocated (num_staging_layers)"));
   }
   return absl::OkStatus();
 }
@@ -453,7 +459,8 @@ KVCacheManagerBase::KVCacheManagerBase(
     HostBufferAllocator host_allocator, std::optional<std::string> bind_ip,
     std::optional<size_t> logical_slice_byte_size,
     std::vector<int64_t> logical_dimensions,
-    std::optional<size_t> logical_physical_size)
+    std::optional<size_t> logical_physical_size,
+    std::optional<int> num_staging_layers)
     : RaidenManagerBase(
           layer_buffers.size(), raiden::DetectShards(layer_buffers),
           logical_slice_byte_size.has_value()
@@ -464,6 +471,10 @@ KVCacheManagerBase::KVCacheManagerBase(
           local_port, parallelism, bind_ip,
           raiden::DetectNumaNodes(layer_buffers), GetLocalHostNicAddresses()),
       host_allocator_(host_allocator) {
+  // Validate before the empty-cache early return so an out-of-range value is
+  // rejected regardless of how many layers were registered.
+  const size_t num_host_layers_to_allocate =
+      ResolveNumStagingLayers(num_layers_, num_staging_layers);
   if (num_layers_ == 0 || num_shards_ == 0) {
     return;
   }
@@ -538,9 +549,8 @@ KVCacheManagerBase::KVCacheManagerBase(
     buffer_holds_.push_back(std::move(device_info));
   }
   absl::StatusOr<HostLayers> host_layers = AllocateHostLayers(
-      num_layers_, num_shards_,
-      /*num_layers_to_allocate=*/ResolveEnvNumStagingLayers(num_layers_),
-      num_host_blocks, buffer_holds_, host_allocator);
+      num_layers_, num_shards_, num_host_layers_to_allocate, num_host_blocks,
+      buffer_holds_, host_allocator);
   if (!host_layers.ok()) {
     // Constructors cannot return a status; a failed host allocation surfaces
     // the way it always has.
@@ -567,17 +577,18 @@ KVCacheManagerBase::KVCacheManagerBase(
     size_t num_layers, size_t num_shards, size_t slice_byte_size,
     std::optional<int> local_port, std::optional<int> host_blocks_to_allocate,
     int parallelism, HostBufferAllocator host_allocator,
-    std::optional<std::string> bind_ip)
+    std::optional<std::string> bind_ip, std::optional<int> num_staging_layers)
     : KVCacheManagerBase(num_layers, num_shards,
                          std::vector<size_t>(num_layers, slice_byte_size),
                          local_port, host_blocks_to_allocate, parallelism,
-                         std::move(host_allocator), std::move(bind_ip)) {}
+                         std::move(host_allocator), std::move(bind_ip),
+                         num_staging_layers) {}
 
 KVCacheManagerBase::KVCacheManagerBase(
     size_t num_layers, size_t num_shards, std::vector<size_t> slice_byte_sizes,
     std::optional<int> local_port, std::optional<int> host_blocks_to_allocate,
     int parallelism, HostBufferAllocator host_allocator,
-    std::optional<std::string> bind_ip)
+    std::optional<std::string> bind_ip, std::optional<int> num_staging_layers)
     : RaidenManagerBase(num_layers, BuildDefaultShards(num_shards),
                         slice_byte_sizes.empty() ? 0 : slice_byte_sizes[0],
                         local_port, parallelism, bind_ip),
@@ -600,7 +611,8 @@ KVCacheManagerBase::KVCacheManagerBase(
 
   absl::StatusOr<HostLayers> host_layers = AllocateHostLayers(
       num_layers_, num_shards_,
-      /*num_layers_to_allocate=*/ResolveEnvNumStagingLayers(num_layers_),
+      /*num_layers_to_allocate=*/
+      ResolveNumStagingLayers(num_layers_, num_staging_layers),
       /*num_host_blocks=*/total_blocks, /*device_layers=*/{}, host_allocator);
   if (!host_layers.ok()) {
     throw std::runtime_error(std::string(host_layers.status().message()));

tpu_sync/kv_cache/kv_cache_manager_base.h

@@ -117,7 +117,9 @@ using HostBlockId = int64_t;
 class KVCacheManagerBase : public tpu_raiden::RaidenManagerBase {
  public:
   // Core C++ Constructor wrapping raw PJRT buffers directly (used by JAX and
-  // PyTorch E2E)
+  // PyTorch E2E). |num_staging_layers| bounds the host mirror along the layer
+  // dimension to that many host layer buffers, in [1, num_layers]; unset keeps
+  // one host buffer per layer. Throws std::invalid_argument when out of range.
   KVCacheManagerBase(
       const std::vector<std::vector<raiden::RaidenBufferHandle>>& layer_buffers,
       std::optional<int> local_port = std::nullopt,
@@ -127,26 +129,30 @@ class KVCacheManagerBase : public tpu_raiden::RaidenManagerBase {
       std::optional<std::string> bind_ip = std::nullopt,
       std::optional<size_t> logical_slice_byte_size = std::nullopt,
       std::vector<int64_t> logical_dimensions = {},
-      std::optional<size_t> logical_physical_size = std::nullopt);
+      std::optional<size_t> logical_physical_size = std::nullopt,
+      std::optional<int> num_staging_layers = std::nullopt);
 
-  // Standard CPU-only constructor for remote workers E2E.
+  // Standard CPU-only constructor for remote workers E2E. |num_staging_layers|
+  // as in the device-buffer constructor.
   KVCacheManagerBase(size_t num_layers, size_t num_shards,
                      size_t slice_byte_size,
                      std::optional<int> local_port = std::nullopt,
                      std::optional<int> host_blocks_to_allocate = std::nullopt,
                      int parallelism = 1,
                      HostBufferAllocator host_allocator = nullptr,
-                     std::optional<std::string> bind_ip = std::nullopt);
+                     std::optional<std::string> bind_ip = std::nullopt,
+                     std::optional<int> num_staging_layers = std::nullopt);
 
   // same as above, but with a vector of slice byte sizes for heterogeneous
-  // registration.
+  // registration. |num_staging_layers| as in the device-buffer constructor.
   KVCacheManagerBase(size_t num_layers, size_t num_shards,
                      std::vector<size_t> slice_byte_sizes,
                      std::optional<int> local_port = std::nullopt,
                      std::optional<int> host_blocks_to_allocate = std::nullopt,
                      int parallelism = 1,
                      HostBufferAllocator host_allocator = nullptr,
-                     std::optional<std::string> bind_ip = std::nullopt);
+                     std::optional<std::string> bind_ip = std::nullopt,
+                     std::optional<int> num_staging_layers = std::nullopt);
 
   ~KVCacheManagerBase() override;
 
@@ -796,7 +802,7 @@ class KVCacheManagerBase : public tpu_raiden::RaidenManagerBase {
                                int block_id) override;
 
   // Returns the number of allocated host layer buffers (equals num_layers()
-  // unless bounded via TPU_RAIDEN_NUM_STAGING_LAYERS).
+  // unless bounded via the num_staging_layers constructor argument).
   size_t num_host_layers() const { return layers_.size(); }
 
   // Returns the total number of bytes allocated in host DRAM.