fix: split page create from map to avoid a remap race on Volta

csrc/allocator.cpp

@@ -256,6 +256,57 @@ bool FTensorAllocator::unmap_from_kv_tensors(
   return true;
 }
 
+bool FTensorAllocator::for_each_mapping_(
+    const std::vector<offset_t> &offsets,
+    const std::function<bool(FTensor *, offset_t)> &op) {
+  if (num_layers_ == 0) {
+    LOGGER(ERROR, "try to map to KV tensors when KV tensors are not created");
+    return false;
+  }
+
+  if (contiguous_layout_) {
+    // Single contiguous tensor; each offset covers all layers.
+    auto ftensor = contiguous_kv_tensor_.get();
+    for (auto offset : offsets) {
+      op(ftensor, offset);
+    }
+  } else if (unified_pool_) {
+    // Unified pool: one block-interleaved FTensor per layer, one page per pid.
+    for (int64_t i = 0; i < num_layers_; i++) {
+      auto kv_name = std::string(kv_prefix) + std::to_string(i);
+      auto ftensor = ftensors_[kv_name].get();
+      for (auto offset : offsets) {
+        op(ftensor, offset);
+      }
+    }
+  } else {
+    // Original per-layer layout: K and V are stacked at the 1st dim.
+    for (int64_t i = 0; i < num_layers_; i++) {
+      auto kv_name = std::string(kv_prefix) + std::to_string(i);
+      auto ftensor = ftensors_[kv_name].get();
+      auto v_base_offset = get_v_base_offset(ftensor->get_tensor());
+      for (auto offset : offsets) {
+        op(ftensor, offset);
+        op(ftensor, offset + v_base_offset);
+      }
+    }
+  }
+  return true;
+}
+
+bool FTensorAllocator::prepare_kv_tensors(
+    const std::vector<offset_t> &offsets) {
+  std::unique_lock<std::mutex> lock(mtx_);
+  return for_each_mapping_(
+      offsets, [](FTensor *ft, offset_t off) { return ft->prepare(off); });
+}
+
+bool FTensorAllocator::commit_kv_tensors(const std::vector<offset_t> &offsets) {
+  std::unique_lock<std::mutex> lock(mtx_);
+  return for_each_mapping_(
+      offsets, [](FTensor *ft, offset_t off) { return ft->commit(off); });
+}
+
 std::string FTensorAllocator::get_anon_tensor_name_() {
   static constexpr std::string_view prefix = "anon_tensor_";
   static std::atomic<int> counter(0);

csrc/ftensor.cpp

@@ -93,34 +93,75 @@ FTensor::~FTensor() {
       ASSERT(munmap(vaddr_, size_) == 0, "munmap failed.");
     }
   }
-  mapping_.clear(); // Free physical page handles after their mappings are gone.
+  mapping_.clear();  // Free physical page handles after their mappings are gone.
+  prepared_.clear(); // Free prepared-but-uncommitted handles (no VA to undo).
   zero_page_.reset();
 }
 
-bool FTensor::map(offset_t offset) {
+bool FTensor::map(offset_t offset) { return prepare(offset) && commit(offset); }
+
+bool FTensor::prepare(offset_t offset) {
   assert(offset % page_size_ == 0); // Ensure alignment.
 
   page_id_t page_id = offset / page_size_;
   if (mapping_.find(page_id) != mapping_.end()) {
     LOGGER(ERROR, "Page %ld is already mapped.", page_id);
     return false;
   }
+  if (prepared_.find(page_id) != prepared_.end()) {
+    return true; // Already prepared; nothing to do.
+  }
+
+  // Physical page allocation only (cuMemCreate). Touches no mapped VA, so this
+  // is safe to run on the background prealloc thread while kernels are running.
+  prepared_[page_id] = make_unique_page(dev_, page_id, page_size_);
+  return true;
+}
+
+bool FTensor::commit(offset_t offset) {
+  assert(offset % page_size_ == 0); // Ensure alignment.
+
+  page_id_t page_id = offset / page_size_;
+  if (mapping_.find(page_id) != mapping_.end()) {
+    // Already committed. Happens when a freed page is kept in the reserved pool
+    // while still mapped (free fast-path) and then handed back out: no VA edit
+    // is needed. commit() is idempotent for such pages.
+    return true;
+  }
+  auto it = prepared_.find(page_id);
+  if (it == prepared_.end()) {
+    // No prior prepare() (on-demand path): create the page now.
+    if (!prepare(offset)) {
+      return false;
+    }
+    it = prepared_.find(page_id);
+  }
 
   auto vaddr = reinterpret_cast<generic_ptr_t>(
       reinterpret_cast<uintptr_t>(vaddr_) + offset);
+  // VA page-table edit: unmap the shared zero page, then map the real page.
+  // The caller must ensure the GPU is idle here (device_synchronize) so this
+  // edit never overlaps an in-flight kernel.
   if (dev_.is_cuda()) {
     CHECK_GPU(gpu_vmm::mem_unmap(vaddr, page_size_));
   }
-
-  mapping_[page_id] = make_unique_page(dev_, page_id, page_size_);
-  mapping_[page_id]->map(vaddr);
+  it->second->map(vaddr);
+  mapping_[page_id] = std::move(it->second);
+  prepared_.erase(it);
   return true;
 }
 
 bool FTensor::unmap(offset_t offset) {
   assert(offset % page_size_ == 0); // Ensure alignment.
 
   page_id_t page_id = offset / page_size_;
+  // A page prepared but never committed still has its VA on the zero page;
+  // just drop the physical handle, no VA edit needed.
+  auto pit = prepared_.find(page_id);
+  if (pit != prepared_.end()) {
+    prepared_.erase(pit);
+    return true;
+  }
   if (mapping_.find(page_id) == mapping_.end()) {
     LOGGER(ERROR, "Page %ld is not mapped.", page_id);
     return false;

csrc/inc/allocator.hpp

@@ -4,6 +4,7 @@
 #pragma once
 
 #include <cstddef>
+#include <functional>
 #include <memory>
 #include <mutex>
 #include <string>
@@ -34,6 +35,11 @@ class FTensorAllocator {
   bool kv_tensors_created();
   bool map_to_kv_tensors(const std::vector<offset_t> &offsets);
   bool unmap_from_kv_tensors(const std::vector<offset_t> &offsets);
+  // Split of map_to_kv_tensors(): prepare_ allocates physical pages only (safe
+  // off the main thread); commit_ does the VA edits (must be GPU-idle). See
+  // FTensor::prepare/commit.
+  bool prepare_kv_tensors(const std::vector<offset_t> &offsets);
+  bool commit_kv_tensors(const std::vector<offset_t> &offsets);
 
   // Global status interfaces.
   // init() creates the default allocator (group_id=0).
@@ -59,6 +65,10 @@ class FTensorAllocator {
   at::Tensor create_ftensor_(size_t size, c10::ScalarType dtype,
                              const std::string &dev_str, std::string name = "");
   void free_ftensor_(at::Tensor &ftensor);
+  // Apply `op` (FTensor::prepare/commit) to every (ftensor, offset) slot
+  // implied by the current layout. Must be called with mtx_ held.
+  bool for_each_mapping_(const std::vector<offset_t> &offsets,
+                         const std::function<bool(FTensor *, offset_t)> &op);
 
   // GPU VMM util functions.
   void init_gpu_();

csrc/inc/ftensor.hpp

@@ -23,6 +23,14 @@ class FTensor {
           size_t page_size = 0);
   ~FTensor();
   bool map(offset_t offset);
+  // Split of map() into a create half and a VA-edit half, so the two can run
+  // on different threads. prepare() does only the physical allocation
+  // (cuMemCreate) and stashes the page; it touches no mapped VA and is safe to
+  // run concurrently with in-flight kernels. commit() does the VA page-table
+  // edit (unmap zero page -> map the prepared page) and must run at a GPU-idle
+  // point. map() = prepare() + commit().
+  bool prepare(offset_t offset);
+  bool commit(offset_t offset);
   bool unmap(offset_t offset);
 
   inline at::Tensor get_tensor() noexcept { return tensor_; }
@@ -42,6 +50,10 @@ class FTensor {
 
   at::Tensor tensor_;
   std::unordered_map<page_id_t, std::unique_ptr<Page>> mapping_;
+  // Pages created by prepare() but not yet committed (VA still on the zero
+  // page). commit() moves an entry from here to mapping_; unmap() discards one
+  // from here if the page was never committed.
+  std::unordered_map<page_id_t, std::unique_ptr<Page>> prepared_;
 };
 
 } // namespace kvcached

csrc/inc/page_allocator.hpp

@@ -134,8 +134,15 @@ class PageAllocator {
   void resize_watcher();
 
   // Internal methods
-  void map_pages(const std::vector<page_id_t> &page_ids);
+  // Page mapping is split so the expensive physical allocation can run on the
+  // background prealloc thread while the cheap VA edit runs on the main thread
+  // at a GPU-idle point. prepare_pages() allocates; commit_pages() maps.
+  void prepare_pages(const std::vector<page_id_t> &page_ids);
+  void commit_pages(const std::vector<page_id_t> &page_ids);
   void unmap_pages(const std::vector<page_id_t> &page_ids);
+  std::vector<offset_t> page_offsets_(const std::vector<page_id_t> &page_ids)
+      const;
+  bool uses_broadcast_() const;
   int64_t get_num_inuse_pages_unlocked() const;
   PageState get_page_state_unlocked() const;
   void update_memory_usage_unlocked();

csrc/page_allocator.cpp

@@ -170,30 +170,24 @@ std::shared_ptr<InternalPage> PageAllocator::alloc_page() {
 
   std::unique_lock<std::mutex> lock(lock_);
   page_id_t page_id = -1;
+  bool from_reserved = false;
 
   while (page_id == -1) {
-    // Fast path: allocate from reserved pages
+    // Fast path: allocate from reserved pages (already physically allocated by
+    // the background prealloc thread; only the VA edit remains).
     if (!reserved_page_list_.empty()) {
       page_id = reserved_page_list_.front();
       reserved_page_list_.pop_front();
       num_free_pages_.fetch_sub(1, std::memory_order_relaxed);
+      from_reserved = true;
 
       // Trigger preallocation to refill reserved pool if getting low
       if (reserved_page_list_.size() <
           static_cast<size_t>(min_reserved_pages_)) {
         prealloc_needed_ = true;
         cond_.notify_all();
       }
-
-      update_memory_usage_unlocked();
-      auto end_time = std::chrono::steady_clock::now();
-      auto duration = std::chrono::duration_cast<std::chrono::microseconds>(
-          end_time - start_time);
-      LOGGER(DEBUG, "alloc 1 page fast path cost %lu us", duration.count());
-      // std::cout << "alloc 1 page fast path cost " << duration.count() << "
-      // us" << std::endl;
-
-      return std::make_shared<InternalPage>(page_id, page_size_);
+      break;
     }
 
     // Slow path: allocate from free pages
@@ -220,7 +214,13 @@ std::shared_ptr<InternalPage> PageAllocator::alloc_page() {
   lock.unlock();
 
   try {
-    map_pages({page_id});
+    // A reserved page was already prepared (cuMemCreate) in the background; a
+    // free page has not been, so prepare it here. The VA edit (commit) always
+    // runs on this main thread at a GPU-idle point, so it never races a kernel.
+    if (!from_reserved) {
+      prepare_pages({page_id});
+    }
+    commit_pages({page_id});
   } catch (const std::exception &e) {
     std::lock_guard<std::mutex> guard(lock_);
     free_page_list_.push_front(page_id);
@@ -230,7 +230,9 @@ std::shared_ptr<InternalPage> PageAllocator::alloc_page() {
                              ": " + e.what());
   }
 
-  if (enable_page_prealloc_) {
+  // Refill the reserved pool after draining a free page on the slow path; the
+  // fast path already signalled above when the pool ran low.
+  if (enable_page_prealloc_ && !from_reserved) {
     trigger_preallocation();
   }
 
@@ -640,7 +642,9 @@ void PageAllocator::prealloc_worker() {
 
     if (!pages_to_reserve.empty()) {
       try {
-        map_pages(pages_to_reserve);
+        // Background thread does the physical allocation only; the VA edit is
+        // deferred to commit_pages() on the main thread (see alloc_page).
+        prepare_pages(pages_to_reserve);
         lock.lock();
         reserved_page_list_.insert(reserved_page_list_.end(),
                                    pages_to_reserve.begin(),
@@ -668,50 +672,68 @@ void PageAllocator::prealloc_worker() {
   }
 }
 
-void PageAllocator::map_pages(const std::vector<page_id_t> &page_ids) {
+std::vector<offset_t>
+PageAllocator::page_offsets_(const std::vector<page_id_t> &page_ids) const {
   std::vector<offset_t> offsets;
   offsets.reserve(page_ids.size());
-
-  if (contiguous_layout_) {
-    for (page_id_t pid : page_ids) {
-      offsets.push_back(pid * page_size_ * num_layers_ * num_kv_buffers_);
-    }
-  } else {
-    for (page_id_t pid : page_ids) {
-      offsets.push_back(pid * page_size_);
-    }
+  const int64_t stride =
+      contiguous_layout_ ? page_size_ * num_layers_ * num_kv_buffers_
+                         : page_size_;
+  for (page_id_t pid : page_ids) {
+    offsets.push_back(pid * stride);
   }
+  return offsets;
+}
+
+bool PageAllocator::uses_broadcast_() const {
+  return (world_size_ > 1 || should_use_worker_ipc()) && broadcast_map_callback_;
+}
 
-  if ((world_size_ > 1 || should_use_worker_ipc()) && broadcast_map_callback_) {
-    // Multi-process mode: execute map on all TP workers via broadcast callback
+void PageAllocator::prepare_pages(const std::vector<page_id_t> &page_ids) {
+  auto offsets = page_offsets_(page_ids);
+
+  if (uses_broadcast_()) {
+    // Multi-process mode has no prepare/commit split over IPC, so do the full
+    // map here (commit_pages() is then a no-op). This still maps from whatever
+    // thread calls prepare_pages(); the create/map split is single-process
+    // only for now.
     broadcast_map_callback_(world_size_, offsets);
   } else {
-    // Single-process mode: directly call FTensorAllocator
+    // Single-process mode: physical allocation only (no VA edit yet).
     auto allocator = FTensorAllocator::global_allocator(group_id_);
-    bool success = allocator->map_to_kv_tensors(offsets);
-    if (!success) {
-      throw std::runtime_error("Failed to map pages to KV tensors");
+    if (!allocator->prepare_kv_tensors(offsets)) {
+      throw std::runtime_error("Failed to prepare pages for KV tensors");
     }
   }
 
-  LOGGER(INFO, "Mapped %zu pages to KV tensors", page_ids.size());
+  LOGGER(INFO, "Prepared %zu pages for KV tensors", page_ids.size());
+}
+
+void PageAllocator::commit_pages(const std::vector<page_id_t> &page_ids) {
+  if (uses_broadcast_()) {
+    // Already mapped in prepare_pages() for the broadcast path.
+    return;
+  }
+
+  auto offsets = page_offsets_(page_ids);
+  // The VA page-table edit must not overlap an in-flight kernel. Callers reach
+  // here on the main thread at a GPU-idle point, but synchronize to guarantee
+  // it in async scheduling mode.
+  if (async_sched_) {
+    CHECK_GPU(gpu_vmm::device_synchronize());
+  }
+  auto allocator = FTensorAllocator::global_allocator(group_id_);
+  if (!allocator->commit_kv_tensors(offsets)) {
+    throw std::runtime_error("Failed to commit pages to KV tensors");
+  }
+
+  LOGGER(INFO, "Committed %zu pages to KV tensors", page_ids.size());
 }
 
 void PageAllocator::unmap_pages(const std::vector<page_id_t> &page_ids) {
   auto start_time = std::chrono::steady_clock::now();
 
-  std::vector<offset_t> offsets;
-  offsets.reserve(page_ids.size());
-
-  if (contiguous_layout_) {
-    for (page_id_t pid : page_ids) {
-      offsets.push_back(pid * page_size_ * num_layers_ * num_kv_buffers_);
-    }
-  } else {
-    for (page_id_t pid : page_ids) {
-      offsets.push_back(pid * page_size_);
-    }
-  }
+  auto offsets = page_offsets_(page_ids);
 
   if ((world_size_ > 1 || should_use_worker_ipc()) &&
       broadcast_unmap_callback_) {