diff --git a/csrc/allocator.cpp b/csrc/allocator.cpp index edb4f34c..42109a20 100644 --- a/csrc/allocator.cpp +++ b/csrc/allocator.cpp @@ -256,6 +256,57 @@ bool FTensorAllocator::unmap_from_kv_tensors( return true; } +bool FTensorAllocator::for_each_mapping_( + const std::vector &offsets, + const std::function &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 &offsets) { + std::unique_lock lock(mtx_); + return for_each_mapping_( + offsets, [](FTensor *ft, offset_t off) { return ft->prepare(off); }); +} + +bool FTensorAllocator::commit_kv_tensors(const std::vector &offsets) { + std::unique_lock 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 counter(0); diff --git a/csrc/ftensor.cpp b/csrc/ftensor.cpp index 95ad9f23..a2c98bbd 100644 --- a/csrc/ftensor.cpp +++ b/csrc/ftensor.cpp @@ -93,11 +93,14 @@ 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_; @@ -105,15 +108,46 @@ bool FTensor::map(offset_t offset) { 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( reinterpret_cast(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; } @@ -121,6 +155,13 @@ 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; diff --git a/csrc/inc/allocator.hpp b/csrc/inc/allocator.hpp index c989c52f..ef173168 100644 --- a/csrc/inc/allocator.hpp +++ b/csrc/inc/allocator.hpp @@ -4,6 +4,7 @@ #pragma once #include +#include #include #include #include @@ -34,6 +35,11 @@ class FTensorAllocator { bool kv_tensors_created(); bool map_to_kv_tensors(const std::vector &offsets); bool unmap_from_kv_tensors(const std::vector &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 &offsets); + bool commit_kv_tensors(const std::vector &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 &offsets, + const std::function &op); // GPU VMM util functions. void init_gpu_(); diff --git a/csrc/inc/ftensor.hpp b/csrc/inc/ftensor.hpp index b56d19b6..f0328f6d 100644 --- a/csrc/inc/ftensor.hpp +++ b/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> 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> prepared_; }; } // namespace kvcached diff --git a/csrc/inc/page_allocator.hpp b/csrc/inc/page_allocator.hpp index 697f5256..c327e7fa 100644 --- a/csrc/inc/page_allocator.hpp +++ b/csrc/inc/page_allocator.hpp @@ -134,8 +134,15 @@ class PageAllocator { void resize_watcher(); // Internal methods - void map_pages(const std::vector &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_ids); + void commit_pages(const std::vector &page_ids); void unmap_pages(const std::vector &page_ids); + std::vector page_offsets_(const std::vector &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(); diff --git a/csrc/page_allocator.cpp b/csrc/page_allocator.cpp index 24bda3a2..60bcbeab 100644 --- a/csrc/page_allocator.cpp +++ b/csrc/page_allocator.cpp @@ -170,13 +170,16 @@ std::shared_ptr PageAllocator::alloc_page() { std::unique_lock 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() < @@ -184,16 +187,7 @@ std::shared_ptr PageAllocator::alloc_page() { prealloc_needed_ = true; cond_.notify_all(); } - - update_memory_usage_unlocked(); - auto end_time = std::chrono::steady_clock::now(); - auto duration = std::chrono::duration_cast( - 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(page_id, page_size_); + break; } // Slow path: allocate from free pages @@ -220,7 +214,13 @@ std::shared_ptr 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 guard(lock_); free_page_list_.push_front(page_id); @@ -230,7 +230,9 @@ std::shared_ptr 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_ids) { +std::vector +PageAllocator::page_offsets_(const std::vector &page_ids) const { std::vector 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_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_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_ids) { auto start_time = std::chrono::steady_clock::now(); - std::vector 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_) { diff --git a/tests/test_create_map_split.py b/tests/test_create_map_split.py new file mode 100644 index 00000000..1fb54f7c --- /dev/null +++ b/tests/test_create_map_split.py @@ -0,0 +1,249 @@ +# SPDX-FileCopyrightText: Copyright contributors to the kvcached project +# SPDX-License-Identifier: Apache-2.0 +"""Regression tests for the page create/map split (FTensor prepare()/commit()). + +The split lets the background prealloc thread do only the physical allocation +(cuMemCreate) while the VA edit (mem_unmap zero page + cuMemMap real page) runs +on the main thread at the alloc_page handout. These tests lock in the observable +accounting invariants that the split must preserve, on the contiguous layout +with prealloc enabled (the config the split targets): + + 1. alloc_page maps distinct physical pages -- no zero_page aliasing across + blocks in different physical pages (prepare()+commit() correctness). + 2. free -> realloc round-trips data intact -- exercises handing back a page + that was reserved while still mapped (commit() must be idempotent) as well + as freshly prepared pages. + 3. trim() drops prealloc'd-but-uncommitted reserved pages without error and + the pool stays usable afterwards -- exercises unmap() on a page that only + exists in FTensor::prepared_, never committed to a real VA. + +Needs the compiled extension and a CUDA/HIP device; skipped otherwise (like the +other real-extension GPU tests here). Contiguous layout is forced via env before +kvcached is imported, so run this module in its own process. + +Run: + KVCACHED_CONTIGUOUS_LAYOUT=true python tests/test_create_map_split.py + # or: pytest tests/test_create_map_split.py +""" + +import os +import time + +# Must be set before importing kvcached (utils reads these at import time). +os.environ.setdefault("KVCACHED_CONTIGUOUS_LAYOUT", "true") +os.environ.setdefault("KVCACHED_PAGE_PREALLOC_ENABLED", "true") + +import torch + +try: + import pytest +except ModuleNotFoundError: + # Allow running as plain `python tests/test_create_map_split.py` in + # environments without pytest (e.g. inside the serving image); the + # decorators degrade to no-ops and the __main__ runner drives the checks. + class _NoPytest: + class mark: + @staticmethod + def skipif(*_a, **_k): + return lambda f: f + + @staticmethod + def fixture(*_a, **_k): + return lambda f: f + + pytest = _NoPytest() # type: ignore[assignment] + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), + reason="create/map split test needs a CUDA/HIP device and the compiled " + "kvcached extension", +) + +# ── Config ────────────────────────────────────────────────────────── +PAGE_TOKENS = 16 # tokens per block ("page" in SGLang terms) +HEAD_NUM = 8 +HEAD_DIM = 64 +NUM_LAYERS = 2 +DTYPE = torch.float16 +DEVICE = "cuda:0" +NUM_TOKENS = 65536 # spans many 2 MB physical pages +NUM_PHYS_PAGES = 12 # how many distinct physical pages to exercise +# ──────────────────────────────────────────────────────────────────── + + +def _setup(): + """Init kvcached, create KV tensors, build the manager. Returns + (k_tensors, manager); caller is responsible for shutdown_kvcached().""" + from kvcached.integration.sglang.interfaces import ( + alloc_kv_cache, + get_kv_cache_manager, + init_kvcached, + ) + from kvcached.vmm_ops import kv_tensors_created + + from kvcached.utils import CONTIGUOUS_LAYOUT + assert CONTIGUOUS_LAYOUT, ( + "this test targets the contiguous layout; run with " + "KVCACHED_CONTIGUOUS_LAYOUT=true in its own process") + + torch.cuda.set_device(0) + init_kvcached(async_sched=False) + + k_tensors, _ = alloc_kv_cache( + kvcache_shape=(NUM_TOKENS, HEAD_NUM, HEAD_DIM), + dtype=DTYPE, + device=DEVICE, + num_layers=NUM_LAYERS, + page_size=PAGE_TOKENS, + attention_type="MHA", + kv_layout="NHD", + ) + + t0 = time.time() + while not kv_tensors_created(): + assert time.time() - t0 < 10, "KV tensors not created within 10 s" + time.sleep(0.1) + + cell_size = HEAD_NUM * HEAD_DIM * DTYPE.itemsize + manager = get_kv_cache_manager( + num_blocks=NUM_TOKENS // PAGE_TOKENS + 1, + block_size=PAGE_TOKENS, + cell_size=cell_size, + num_layers=NUM_LAYERS, + reserve_null_block=True, + ) + manager._post_init_done.wait(timeout=10.0) + assert manager._post_init_done.is_set(), "post-init timed out" + return k_tensors, manager + + +@pytest.fixture(scope="module") +def kv_setup(): + """One init/shutdown for the module (kvcached global state is a singleton).""" + from kvcached.integration.sglang.interfaces import shutdown_kvcached + k_tensors, manager = _setup() + yield k_tensors, manager + shutdown_kvcached() + + +def _blocks_per_phys(page_size_bytes): + return page_size_bytes // (PAGE_TOKENS * HEAD_NUM * HEAD_DIM * DTYPE.itemsize) + + +def _write_read_distinct(k_buf, block_ids, blocks_per_phys, base): + """Write base+i to the first token of each physical page, read it back.""" + tokens = [block_ids[i * blocks_per_phys] * PAGE_TOKENS + for i in range(NUM_PHYS_PAGES)] + for i, tok in enumerate(tokens): + k_buf[tok] = torch.full((HEAD_NUM, HEAD_DIM), float(base + i), + dtype=DTYPE, device=DEVICE) + torch.cuda.synchronize() + got = [k_buf[tok][0][0].item() for tok in tokens] + torch.cuda.synchronize() + return got, [float(base + i) for i in range(NUM_PHYS_PAGES)] + + +# NOTE: defined first on purpose. Before any alloc, the reserved pool holds +# only prealloc'd-but-uncommitted pages, which is exactly the prepared-page +# drop path this exercises. Later tests reserve mapped pages too (free +# fast-path), which would dilute it. +def test_trim_drops_uncommitted_reserved_pages(kv_setup): + """trim() must release prealloc'd-but-uncommitted reserved pages (present + only in FTensor::prepared_, never committed to a real VA) without error, + and the pool must stay usable afterwards.""" + k_tensors, manager = kv_setup + from kvcached.utils import PAGE_SIZE + pa = manager.page_allocator + k_buf = k_tensors[0] + bpp = _blocks_per_phys(PAGE_SIZE) + + # Wait for the initial prealloc fill, then stop the thread so the reserved + # count can't change under the asserts (a prealloc already in flight could + # otherwise re-insert just after trim()). + t0 = time.time() + while pa.get_num_reserved_pages() == 0: + assert time.time() - t0 < 10, "prealloc never reserved any pages" + time.sleep(0.05) + pa.stop_prealloc_thread() + + assert pa.get_num_reserved_pages() > 0 # prepared-but-uncommitted pages + pa.trim() # drops them with no VA edit + assert pa.get_num_reserved_pages() == 0 + + pa.start_prealloc_thread() # restore for the remaining tests + + block_ids = manager.alloc(bpp * NUM_PHYS_PAGES) + assert block_ids is not None + try: + got, expected = _write_read_distinct(k_buf, block_ids, bpp, base=200) + assert got == expected, f"corruption after trim: {got} != {expected}" + finally: + manager.free(block_ids) + + +def test_alloc_maps_distinct_physical_pages(kv_setup): + """prepare()+commit() must give each physical page its own mapping.""" + k_tensors, manager = kv_setup + from kvcached.utils import PAGE_SIZE + k_buf = k_tensors[0] + bpp = _blocks_per_phys(PAGE_SIZE) + assert manager.available_size() >= bpp * NUM_PHYS_PAGES + + block_ids = manager.alloc(bpp * NUM_PHYS_PAGES) + assert block_ids is not None + try: + got, expected = _write_read_distinct(k_buf, block_ids, bpp, base=0) + assert got == expected, f"aliasing / bad map: {got} != {expected}" + finally: + manager.free(block_ids) + + +def test_free_realloc_preserves_integrity(kv_setup): + """Handing a page back out -- reserved-while-mapped (idempotent commit) or + freshly prepared -- must still serve correct, non-aliased memory. + + Reliably hits the idempotent-commit branch only with the default + min_reserved < max_reserved, which leaves room for free() to reserve pages + that are still mapped; those then come back out via the fast path. + """ + k_tensors, manager = kv_setup + from kvcached.utils import PAGE_SIZE + k_buf = k_tensors[0] + bpp = _blocks_per_phys(PAGE_SIZE) + + first = manager.alloc(bpp * NUM_PHYS_PAGES) + assert first is not None + manager.free(first) + + second = manager.alloc(bpp * NUM_PHYS_PAGES) + assert second is not None + try: + got, expected = _write_read_distinct(k_buf, second, bpp, base=100) + assert got == expected, f"post-realloc corruption: {got} != {expected}" + finally: + manager.free(second) + + +if __name__ == "__main__": + # Plain-python runner (no pytest needed): share one setup across the checks. + import sys + + from kvcached.integration.sglang.interfaces import shutdown_kvcached + + setup = _setup() + checks = [ + test_trim_drops_uncommitted_reserved_pages, + test_alloc_maps_distinct_physical_pages, + test_free_realloc_preserves_integrity, + ] + failed = 0 + for check in checks: + try: + check(setup) + print(f"[PASS] {check.__name__}") + except Exception as e: # noqa: BLE001 + failed += 1 + print(f"[FAIL] {check.__name__}: {type(e).__name__}: {e}") + shutdown_kvcached() + print(f"\n{'=' * 50}\nResults: {len(checks) - failed} passed, {failed} failed") + sys.exit(1 if failed else 0)