From c6be7ce1f5c22d3bfb0cb384bc92fdbc7d75411e Mon Sep 17 00:00:00 2001 From: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> Date: Fri, 23 Oct 2026 18:38:44 +0000 Subject: [PATCH 1/5] [None][feat] Per-locality-domain KV cache storage, allocation, eviction and resize in the C++ KV cache manager v2 On a GPU with locality domains (Rubin), each domain's KV cache lives in its own physical memory. The C++ KVCacheManagerV2 now keeps one physical chunk pool per domain, mapped at explicit offsets of one VA reservation per pool, so slot ids stay global and base + slotSize * id is unchanged for every consumer. - Allocation, eviction, migration, resize and quota are per domain; a domain's share of the quota is taken from the raw quota on init and resize. - ReuseScope carries the domain, so prefix reuse stays within a domain. - LocalityDomainId is a StrongIndex like SlotId. - A domain chunk must start on a physical-page boundary in every pool; when the slot sizes need an oversized chunk the manager keeps one domain. - Whole-pool consumers get the mapped ranges; the per-domain free-slot count is an introspection hook for tests. - TRT_LLM_MOCK_LOCALIZATION_SUPPORT=1 splits the cache on any GPU for testing. Signed-off-by: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> --- cpp/tensorrt_llm/CMakeLists.txt | 3 +- cpp/tensorrt_llm/batch_manager/CMakeLists.txt | 5 +- .../kv_cache_manager_v2/CMakeLists.txt | 1 + .../kv_cache_manager_v2/blockRadixTree.cpp | 9 + .../kv_cache_manager_v2/blockRadixTree.h | 7 +- .../kv_cache_manager_v2/coldPageCodec.h | 4 + .../kv_cache_manager_v2/common.h | 3 + .../kv_cache_manager_v2/config.h | 2 + .../kv_cache_manager_v2/cudaVirtMem.cpp | 238 ++++++-- .../kv_cache_manager_v2/cudaVirtMem.h | 92 ++- .../evictionController.cpp | 60 +- .../kv_cache_manager_v2/evictionController.h | 15 +- .../kv_cache_manager_v2/kvCache.cpp | 15 +- .../kv_cache_manager_v2/kvCache.h | 5 + .../kv_cache_manager_v2/kvCacheManager.cpp | 39 +- .../kv_cache_manager_v2/kvCacheManager.h | 4 + .../kv_cache_manager_v2/localityDomain.cpp | 70 +++ .../kv_cache_manager_v2/localityDomain.h | 38 ++ .../kv_cache_manager_v2/page.cpp | 9 +- .../kv_cache_manager_v2/storage/core.cpp | 426 ++++++++++++-- .../kv_cache_manager_v2/storage/core.h | 247 +++++++- .../kv_cache_manager_v2/storageManager.cpp | 313 +++++++--- .../kv_cache_manager_v2/storageManager.h | 74 ++- .../batch_manager/kvCacheManagerV2.cpp | 87 ++- cpp/tensorrt_llm/runtime/CMakeLists.txt | 12 +- .../unit_tests/batch_manager/CMakeLists.txt | 13 + .../kvCacheManagerV2LocalityTest.cpp | 536 ++++++++++++++++++ .../kvCacheManagerV2TypedIndexTest.cpp | 10 +- .../runtime/kv_cache_manager_v2/__init__.pyi | 12 + .../kv_cache_manager_v2/_introspection.py | 5 + .../test_kv_cache_localization.py | 302 ++++++++++ .../test_kv_cache_localization_wrapper.py | 176 ++++++ .../test_kv_cache_salting.py | 13 + 33 files changed, 2552 insertions(+), 293 deletions(-) create mode 100644 cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/localityDomain.cpp create mode 100644 cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/localityDomain.h create mode 100644 cpp/tests/unit_tests/batch_manager/kvCacheManagerV2LocalityTest.cpp create mode 100644 tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_localization.py create mode 100644 tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_localization_wrapper.py diff --git a/cpp/tensorrt_llm/CMakeLists.txt b/cpp/tensorrt_llm/CMakeLists.txt index 8fa6c33c8eee..6e44c498d132 100644 --- a/cpp/tensorrt_llm/CMakeLists.txt +++ b/cpp/tensorrt_llm/CMakeLists.txt @@ -273,7 +273,7 @@ target_link_libraries(${SHARED_TARGET} PUBLIC ${TRTLLM_LINK_LIBS}) target_link_libraries( ${SHARED_TARGET} PRIVATE - $ + $ $ $ $ @@ -295,6 +295,7 @@ foreach( tllm_runtime_buffers tllm_runtime_distributed tllm_runtime_topology + tllm_runtime_locality_domain tllm_executor_serialization tllm_common_environment tllm_executor_values diff --git a/cpp/tensorrt_llm/batch_manager/CMakeLists.txt b/cpp/tensorrt_llm/batch_manager/CMakeLists.txt index cdc103605ae2..c6f3639d1832 100644 --- a/cpp/tensorrt_llm/batch_manager/CMakeLists.txt +++ b/cpp/tensorrt_llm/batch_manager/CMakeLists.txt @@ -86,8 +86,9 @@ add_library(tllm::kv_cache ALIAS tllm_kv_cache) target_include_directories( tllm_kv_cache PUBLIC ${KV_CACHE_MANAGER_V2_INCLUDE_DIR} ${SHA256_DIR}) target_link_libraries( - tllm_kv_cache PUBLIC tllm::kv_cache_host_memory tllm::common_cuda - tllm::executor_types CUDA::cuda_driver) + tllm_kv_cache + PUBLIC tllm::kv_cache_host_memory tllm::common_cuda tllm::executor_types + tllm::runtime_locality_domain CUDA::cuda_driver) add_library( tllm_kv_cache_legacy STATIC diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/CMakeLists.txt b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/CMakeLists.txt index abf8ffb9ed82..d07300784ed9 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/CMakeLists.txt +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/CMakeLists.txt @@ -49,6 +49,7 @@ set(KV_CACHE_MANAGER_V2_SRCS kv_cache_manager_v2/common.cpp kv_cache_manager_v2/config.cpp kv_cache_manager_v2/cudaVirtMem.cpp + kv_cache_manager_v2/localityDomain.cpp kv_cache_manager_v2/lifeCycleRegistry.cpp kv_cache_manager_v2/utils/cudaEvent.cpp kv_cache_manager_v2/utils/optionalGilRelease.cpp diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp index 18ed536c27f0..6107c6a83c6f 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp @@ -62,6 +62,10 @@ static void emitReuseScopeBytes(ReuseScope const& scope, Emit&& emit) { mask |= 1U << 1; } + if (scope.localityDomainId.has_value()) + { + mask |= 1U << 2; + } emit(&mask, sizeof(mask)); if (scope.loraId.has_value()) { @@ -74,6 +78,11 @@ static void emitReuseScopeBytes(ReuseScope const& scope, Emit&& emit) emit(reinterpret_cast(&length), sizeof(length)); emit(reinterpret_cast(scope.salt->data()), scope.salt->size()); } + if (scope.localityDomainId.has_value()) + { + std::uint64_t const value = static_cast(scope.localityDomainId->value()); + emit(reinterpret_cast(&value), sizeof(value)); + } } // --------------------------------------------------------------------------- diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h index b6fe8d1dfd4b..c76306ba2f4c 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h @@ -48,16 +48,19 @@ struct Block; // --------------------------------------------------------------------------- // ReuseScope — per-request namespace for prefix reuse. -// Mirrors Python's ReuseScope(lora_id, salt). +// Mirrors Python's ReuseScope(lora_id, salt, locality_domain_id). // --------------------------------------------------------------------------- struct ReuseScope { std::optional loraId; std::optional salt; + // The locality domain is part of the reuse key: pages are backed by one domain's memory, + // so prefixes are only shared between requests on the same domain. + std::optional localityDomainId; bool operator==(ReuseScope const& other) const noexcept { - return loraId == other.loraId && salt == other.salt; + return loraId == other.loraId && salt == other.salt && localityDomainId == other.localityDomainId; } }; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/coldPageCodec.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/coldPageCodec.h index 7e67ecc1e679..8c40e6c977da 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/coldPageCodec.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/coldPageCodec.h @@ -25,6 +25,8 @@ #include #include #include +#include +#include namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 { @@ -35,6 +37,8 @@ struct PoolDesc PoolIndex poolIndex{0}; MemAddress baseAddress = 0; size_t slotBytes = 0; + // Mapped byte ranges [begin, end) relative to baseAddress; a localized pool has holes between them. + std::vector> mappedRanges; }; //! Describes one hot pool group and the lifecycles that share its physical layout. diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.h index 4d867eb906a4..dcb36ac8000a 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.h @@ -114,6 +114,9 @@ inline constexpr FileDescriptor kBadFileDescriptor = -1; using PageIndex = StrongIndex; inline constexpr PageIndex kBadPageIndex{-1}; +// Index of a GPU locality domain. +using LocalityDomainId = StrongIndex; + // Eviction priority (0 = highest priority to evict, 100 = lowest). using Priority = int; inline constexpr Priority kPriorityMin = 0; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h index adb134953938..d5fced290abf 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h @@ -38,6 +38,8 @@ namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 struct GpuCacheTierConfig { size_t quota = 0; // bytes + // Place KV pages in per-locality-domain physical memory when the device exposes locality domains. + bool enableLocalityDomains = false; [[nodiscard]] CacheTier tier() const noexcept { diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/cudaVirtMem.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/cudaVirtMem.cpp index a9aca0a1e185..967bb4f02e18 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/cudaVirtMem.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/cudaVirtMem.cpp @@ -17,10 +17,13 @@ #include "kv_cache_manager_v2/cudaVirtMem.h" #include "kv_cache_manager_v2/exceptions.h" +#include "kv_cache_manager_v2/localityDomain.h" #include "kv_cache_manager_v2/utils/math.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/logger.h" +#include + namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 { @@ -28,12 +31,17 @@ namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 // PooledPhysMemAllocator // --------------------------------------------------------------------------- +static bool isUnsupportedAllocationError(CUresult err) noexcept +{ + return err == CUDA_ERROR_NOT_PERMITTED || err == CUDA_ERROR_NOT_SUPPORTED || err == CUDA_ERROR_INVALID_DEVICE + || err == CUDA_ERROR_INVALID_VALUE; +} + static bool isPropSupported(CUmemAllocationProp const& prop) { CUmemGenericAllocationHandle handle; CUresult err = cuMemCreate(&handle, 2ULL << 20, &prop, 0); - if (err == CUDA_ERROR_NOT_PERMITTED || err == CUDA_ERROR_NOT_SUPPORTED || err == CUDA_ERROR_INVALID_DEVICE - || err == CUDA_ERROR_INVALID_VALUE) + if (isUnsupportedAllocationError(err)) { return false; } @@ -63,22 +71,42 @@ PhysMemWrapper::~PhysMemWrapper() // PooledPhysMemAllocator // --------------------------------------------------------------------------- -PooledPhysMemAllocator::PooledPhysMemAllocator(size_t physMemSize) +PooledPhysMemAllocator::PooledPhysMemAllocator(size_t physMemSize, int numLocalityDomains) : mPhysMemSize(physMemSize) - , mPool([this]() -> PhysMemWrapper* { return new PhysMemWrapper(mPhysMemSize, mProp); }, - [](PhysMemWrapper* chunk) { delete chunk; }) { + TLLM_CHECK_WITH_INFO(numLocalityDomains >= 1, "PooledPhysMemAllocator: need at least one locality domain"); + // Get current device. cuCheck(cuCtxGetDevice(&mDeviceId)); - // Build the best allocation property. Prefer shareable handle types so - // UCX cuda_ipc can export the memory for intra-node zero-copy transfers: - // FABRIC (MNNVL) > POSIX_FILE_DESCRIPTOR (pidfd-based, UCX >= 1.22) > NONE, - // and within each handle type prefer gpuDirectRDMACapable over not. mProp.type = CU_MEM_ALLOCATION_TYPE_PINNED; mProp.location.type = CU_MEM_LOCATION_TYPE_DEVICE; mProp.location.id = mDeviceId; + if (numLocalityDomains > 1 && localizationMode() != LocalizationMode::MOCK) + { + selectLocalizedProp(numLocalityDomains); + } + else + { + selectDefaultProp(); + } + + mPools.reserve(LocalityDomainId{numLocalityDomains}); + for (LocalityDomainId localityDomainId{0}; localityDomainId < numLocalityDomains; ++localityDomainId) + { + mPools.push_back(std::make_unique>([this, localityDomainId]() -> PhysMemWrapper* + { return createChunk(localityDomainId); }, + [](PhysMemWrapper* chunk) { delete chunk; })); + } +} + +void PooledPhysMemAllocator::selectDefaultProp() +{ + // Build the best allocation property. Prefer shareable handle types so + // UCX cuda_ipc can export the memory for intra-node zero-copy transfers: + // FABRIC (MNNVL) > POSIX_FILE_DESCRIPTOR (pidfd-based, UCX >= 1.22) > NONE, + // and within each handle type prefer gpuDirectRDMACapable over not. constexpr CUmemAllocationHandleType kHandleTypePreference[] = {CU_MEM_HANDLE_TYPE_FABRIC, CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR, CU_MEM_HANDLE_TYPE_NONE}; bool supported = false; @@ -114,14 +142,101 @@ PooledPhysMemAllocator::PooledPhysMemAllocator(size_t physMemSize) } } +// Same preference order as the default property, but every domain has to accept the candidate and +// the chunk size is rounded up to a multiple of every domain's granularity. +void PooledPhysMemAllocator::selectLocalizedProp(int numLocalityDomains) +{ + mLocalization.emplace(); + if (!mLocalization->supportsMemoryLocalization()) + { + throw std::runtime_error("PooledPhysMemAllocator: locality domain memory localization is not supported"); + } + + constexpr CUmemAllocationHandleType kHandleTypePreference[] + = {CU_MEM_HANDLE_TYPE_FABRIC, CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR, CU_MEM_HANDLE_TYPE_NONE}; + for (auto const handleType : kHandleTypePreference) + { + for (bool const rdmaCapable : {true, false}) + { + mProp.requestedHandleTypes = handleType; + mProp.allocFlags.gpuDirectRDMACapable = rdmaCapable ? 1 : 0; + + size_t granularity = 1; + bool supported = true; + for (LocalityDomainId localityDomainId{0}; localityDomainId < numLocalityDomains && supported; + ++localityDomainId) + { + size_t domainGranularity = 0; + CUresult const err = mLocalization->tryGetLocalizedAllocationGranularity( + &domainGranularity, localityDomainId.value(), handleType, rdmaCapable, mProp.allocFlags.usage); + if (isUnsupportedAllocationError(err)) + { + supported = false; + break; + } + cuCheck(err); + TLLM_CHECK_WITH_INFO(domainGranularity > 0, "CUDA returned an invalid allocation granularity"); + granularity = std::lcm(granularity, domainGranularity); + } + if (!supported) + { + continue; + } + + size_t const probeSize = roundUp(size_t{2} << 20, granularity); + for (LocalityDomainId localityDomainId{0}; localityDomainId < numLocalityDomains && supported; + ++localityDomainId) + { + CUmemGenericAllocationHandle handle{}; + CUresult const err = mLocalization->tryCreateLocalizedAllocationHandle( + &handle, probeSize, localityDomainId.value(), handleType, rdmaCapable, mProp.allocFlags.usage); + if (isUnsupportedAllocationError(err)) + { + supported = false; + break; + } + cuCheck(err); + cuMemRelease(handle); + } + if (supported) + { + mPhysMemSize = roundUp(mPhysMemSize, granularity); + return; + } + } + } + throw std::runtime_error( + "PooledPhysMemAllocator: no physical memory allocation property is supported by every locality domain"); +} + +PhysMemWrapper* PooledPhysMemAllocator::createChunk(LocalityDomainId localityDomainId) +{ + if (!mLocalization.has_value()) + { + return new PhysMemWrapper(mPhysMemSize, mProp); + } + CUmemGenericAllocationHandle handle{}; + cuCheck(mLocalization->tryCreateLocalizedAllocationHandle(&handle, mPhysMemSize, localityDomainId.value(), + mProp.requestedHandleTypes, mProp.allocFlags.gpuDirectRDMACapable != 0, mProp.allocFlags.usage)); + return new PhysMemWrapper(handle); +} + PooledPhysMemAllocator::~PooledPhysMemAllocator() { - mPool.clear(); + clear(); +} + +PooledPhysMemAllocator::PooledPhysMem PooledPhysMemAllocator::acquire(LocalityDomainId localityDomainId) +{ + return mPools.at(localityDomainId)->get(); } -PooledPhysMemAllocator::PooledPhysMem PooledPhysMemAllocator::acquire() +void PooledPhysMemAllocator::clear() { - return mPool.get(); + for (auto& pool : mPools) + { + pool->clear(); + } } // --------------------------------------------------------------------------- @@ -131,6 +246,7 @@ PooledPhysMemAllocator::PooledPhysMem PooledPhysMemAllocator::acquire() VirtMem::VirtMem(size_t vmSize, PooledPhysMemAllocator& physMemAllocator, size_t initNumPhysMem) : mVmSize(vmSize) , mPhysMemAllocator(physMemAllocator) + , mMappings(LocalityDomainId{physMemAllocator.numLocalityDomains()}) { TLLM_CHECK_DEBUG(vmSize % physMemAllocator.physMemSize() == 0); @@ -142,7 +258,7 @@ VirtMem::VirtMem(size_t vmSize, PooledPhysMemAllocator& physMemAllocator, size_t try { - extend(initNumPhysMem); + extend(initNumPhysMem, LocalityDomainId{0}, size_t{0}); } catch (...) { @@ -159,71 +275,114 @@ VirtMem::~VirtMem() noexcept KVCM2_POISON_ON_EXCEPT([this]() { destroy(); }); } -void VirtMem::push(PooledPhysMemAllocator::PooledPhysMem handle) +void VirtMem::push(PooledPhysMemAllocator::PooledPhysMem handle, size_t offset, LocalityDomainId localityDomainId) { size_t physSize = mPhysMemAllocator.physMemSize(); - CUdeviceptr offset = mAddr + physSize * static_cast(mPhysHandles.size()); - TLLM_CHECK_DEBUG(physSize * (mPhysHandles.size() + 1) <= mVmSize); + CUdeviceptr ptr = mAddr + offset; + TLLM_CHECK_DEBUG(offset + physSize <= mVmSize); - cuCheck(cuMemMap(offset, physSize, 0, handle->handle(), 0)); - cuCheck(cuMemSetAccess(offset, physSize, &mAccessDesc, 1)); - mPhysHandles.push_back(std::move(handle)); + cuCheck(cuMemMap(ptr, physSize, 0, handle->handle(), 0)); + cuCheck(cuMemSetAccess(ptr, physSize, &mAccessDesc, 1)); + mMappings[localityDomainId].push_back(Mapping{offset, std::move(handle)}); + mMappedOffsets.insert(offset); } -void VirtMem::pop() +void VirtMem::pop(LocalityDomainId localityDomainId) { - TLLM_CHECK_DEBUG(!mPhysHandles.empty()); + auto& mappings = mMappings[localityDomainId]; + TLLM_CHECK_DEBUG(!mappings.empty()); size_t physSize = mPhysMemAllocator.physMemSize(); - CUdeviceptr offset = mAddr + physSize * (mPhysHandles.size() - 1); - cuCheck(cuMemUnmap(offset, physSize)); - mPhysHandles.pop_back(); // PhysMemHandle destructor returns to pool + size_t const offset = mappings.back().offset; + cuCheck(cuMemUnmap(mAddr + offset, physSize)); + mMappedOffsets.erase(offset); + mappings.pop_back(); // PhysMemHandle destructor returns to pool } -void VirtMem::extend(size_t numToAdd) +void VirtMem::extend(size_t numToAdd, LocalityDomainId localityDomainId, std::optional offset) { - size_t old = numPhysMem(); + TLLM_CHECK_WITH_INFO(localityDomainId >= 0 && localityDomainId < numLocalityDomains(), + "VirtMem::extend: locality domain %d out of range", localityDomainId.value()); + size_t const physSize = mPhysMemAllocator.physMemSize(); + if (!offset.has_value()) + { + TLLM_CHECK_WITH_INFO(numLocalityDomains() == 1, "VirtMem::extend: offset is required with locality domains"); + offset = mappedBytes(); + } + size_t const begin = *offset; + size_t const end = begin + numToAdd * physSize; + if (begin % physSize != 0 || end > mVmSize) + { + throw LogicError("VirtMem::extend: mapping is outside or misaligned with the reservation"); + } + for (size_t pos = begin; pos < end; pos += physSize) + { + if (mMappedOffsets.count(pos) != 0) + { + throw LogicError("VirtMem::extend: mapping overlaps an existing mapping"); + } + } + + size_t const old = numPhysMem(localityDomainId); try { - for (size_t i = 0; i < numToAdd; ++i) + for (size_t pos = begin; pos < end; pos += physSize) { - push(mPhysMemAllocator.acquire()); + push(mPhysMemAllocator.acquire(localityDomainId), pos, localityDomainId); } } catch (...) { // Rollback: remove any newly-mapped chunks. - while (numPhysMem() > old) + while (numPhysMem(localityDomainId) > old) { - pop(); + pop(localityDomainId); } throw; } } -void VirtMem::shrink(size_t numToRemove) +void VirtMem::shrink(size_t numToRemove, LocalityDomainId localityDomainId) { cuCheck(cuCtxSynchronize()); for (size_t i = 0; i < numToRemove; ++i) { - pop(); + pop(localityDomainId); } } -void VirtMem::realloc(size_t numBytes) +void VirtMem::realloc(size_t numBytes, LocalityDomainId localityDomainId) { size_t physSize = mPhysMemAllocator.physMemSize(); size_t required = divUp(numBytes, physSize); - size_t current = numPhysMem(); + size_t current = numPhysMem(localityDomainId); if (required > current) { - extend(required - current); + extend(required - current, localityDomainId); } else if (required < current) { - shrink(current - required); + shrink(current - required, localityDomainId); } } +std::vector> VirtMem::mappedRanges() const +{ + size_t const physSize = mPhysMemAllocator.physMemSize(); + std::vector> ranges; + for (size_t const offset : mMappedOffsets) + { + if (!ranges.empty() && ranges.back().second == offset) + { + ranges.back().second = offset + physSize; + } + else + { + ranges.emplace_back(offset, offset + physSize); + } + } + return ranges; +} + void VirtMem::destroy() { if (mVmSize == 0) @@ -231,9 +390,12 @@ void VirtMem::destroy() return; } cuCheck(cuCtxSynchronize()); - while (!mPhysHandles.empty()) + for (LocalityDomainId localityDomainId{0}; localityDomainId < numLocalityDomains(); ++localityDomainId) { - pop(); + while (numPhysMem(localityDomainId) != 0) + { + pop(localityDomainId); + } } cuCheck(cuMemAddressFree(mAddr, mVmSize)); mAddr = 0; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/cudaVirtMem.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/cudaVirtMem.h index 5c560f036594..ae486267e51a 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/cudaVirtMem.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/cudaVirtMem.h @@ -19,8 +19,13 @@ #include "kv_cache_manager_v2/common.h" #include "kv_cache_manager_v2/utils/cudaEvent.h" +#include "tensorrt_llm/runtime/locality_domain/locality_domain_utils.h" #include +#include +#include +#include +#include #include namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 @@ -34,6 +39,13 @@ class PhysMemWrapper { public: PhysMemWrapper(size_t size, CUmemAllocationProp const& prop); + + // Adopts a handle created elsewhere (locality-domain allocations). + explicit PhysMemWrapper(CUmemGenericAllocationHandle handle) noexcept + : mHandle(handle) + { + } + ~PhysMemWrapper(); PhysMemWrapper(PhysMemWrapper const&) = delete; @@ -51,6 +63,10 @@ class PhysMemWrapper // --------------------------------------------------------------------------- // PooledPhysMemAllocator — creates and pools physical GPU memory chunks. // Mirrors _cuda_virt_mem.py::PooledPhysMemAllocator. +// +// With more than one locality domain, each domain has its own pool of chunks located in that +// domain's memory (or ordinary device memory in mock mode) and physMemSize() is rounded up to a +// multiple of the localized allocation granularity. // --------------------------------------------------------------------------- class PooledPhysMemAllocator { @@ -58,7 +74,7 @@ class PooledPhysMemAllocator using PooledPhysMem = SimplePool::PoolItem; // physMemSize: size of each physical chunk in bytes. - explicit PooledPhysMemAllocator(size_t physMemSize); + explicit PooledPhysMemAllocator(size_t physMemSize, int numLocalityDomains = 1); ~PooledPhysMemAllocator(); PooledPhysMemAllocator(PooledPhysMemAllocator const&) = delete; @@ -66,14 +82,11 @@ class PooledPhysMemAllocator // Borrow a physical memory handle from the pool (or allocate a new one). // Dropping the returned PhysMemHandle returns it to the pool. - [[nodiscard]] PooledPhysMem acquire(); + [[nodiscard]] PooledPhysMem acquire(LocalityDomainId localityDomainId = LocalityDomainId{0}); // Release all cached (unused) physical memory back to the driver. // Mirrors Python PooledPhysMemAllocator.clear(). - void clear() - { - mPool.clear(); - } + void clear(); [[nodiscard]] size_t physMemSize() const noexcept { @@ -85,20 +98,34 @@ class PooledPhysMemAllocator return mDeviceId; } + int numLocalityDomains() const noexcept + { + return mPools.size().value(); + } + private: + void selectDefaultProp(); + void selectLocalizedProp(int numLocalityDomains); + PhysMemWrapper* createChunk(LocalityDomainId localityDomainId); + size_t mPhysMemSize{}; int mDeviceId{}; CUmemAllocationProp mProp{}; - SimplePool mPool; + // Present only when chunks are really placed per domain; mock mode leaves it empty. + std::optional mLocalization; + TypedVec>> mPools; }; // --------------------------------------------------------------------------- // VirtMem — a virtual address range backed by physical GPU memory chunks. -// Physical chunks are mapped/unmapped at the end of the range (stack discipline). // Mirrors _cuda_virt_mem.py::VirtMem. // -// Invariant: mappedBytes() == numPhysMem() * physMemSize() -// mappedBytes() <= virtualBytes() +// Each locality domain keeps its own stack of mappings. With one domain chunks are mapped and +// unmapped at the end of the range; with several the caller places every chunk at an explicit +// byte offset so the pools sharing a slot-id space agree on which domain backs which slots. +// +// Invariant: mappedBytes(d) == numPhysMem(d) * physMemSize() for every domain d +// sum of mappedBytes(d) <= virtualBytes() // --------------------------------------------------------------------------- class VirtMem { @@ -106,21 +133,23 @@ class VirtMem // vmSize: total virtual address space reserved (bytes). // Must be a multiple of physMemSize. // physMemAllocator: shared allocator for physical chunks. - // initNumPhysMem: number of physical chunks to map immediately. + // initNumPhysMem: number of physical chunks to map immediately (domain 0). VirtMem(size_t vmSize, PooledPhysMemAllocator& physMemAllocator, size_t initNumPhysMem = 0); ~VirtMem() noexcept; VirtMem(VirtMem const&) = delete; VirtMem& operator=(VirtMem const&) = delete; - // Map numPhysMem additional chunks at the top of the address range. - void extend(size_t numPhysMem); + // Map numPhysMem additional chunks of localityDomainId starting at byte offset `offset`. + // offset defaults to the end of the mapped range and must be given with several domains. + void extend(size_t numPhysMem, LocalityDomainId localityDomainId = LocalityDomainId{0}, + std::optional offset = std::nullopt); - // Unmap numPhysMem chunks from the top (synchronizes CUDA first). - void shrink(size_t numPhysMem); + // Unmap the most recently mapped numPhysMem chunks of localityDomainId (synchronizes CUDA first). + void shrink(size_t numPhysMem, LocalityDomainId localityDomainId = LocalityDomainId{0}); - // Adjust mapped bytes to exactly numBytes (extend or shrink). - void realloc(size_t numBytes); + // Adjust mapped bytes of localityDomainId to exactly numBytes (extend or shrink). + void realloc(size_t numBytes, LocalityDomainId localityDomainId = LocalityDomainId{0}); void destroy(); @@ -134,9 +163,9 @@ class VirtMem return mPhysMemAllocator.physMemSize(); } - [[nodiscard]] size_t mappedBytes() const noexcept + [[nodiscard]] size_t mappedBytes(LocalityDomainId localityDomainId = LocalityDomainId{0}) const noexcept { - return mPhysMemAllocator.physMemSize() * numPhysMem(); + return mPhysMemAllocator.physMemSize() * numPhysMem(localityDomainId); } [[nodiscard]] size_t virtualBytes() const noexcept @@ -144,19 +173,34 @@ class VirtMem return mVmSize; } - [[nodiscard]] size_t numPhysMem() const noexcept + [[nodiscard]] size_t numPhysMem(LocalityDomainId localityDomainId = LocalityDomainId{0}) const noexcept { - return mPhysHandles.size(); + return mMappings[localityDomainId].size(); } + [[nodiscard]] int numLocalityDomains() const noexcept + { + return mMappings.size().value(); + } + + // Merged mapped byte ranges [begin, end) relative to address(), across all domains. + [[nodiscard]] std::vector> mappedRanges() const; + private: - void push(PooledPhysMemAllocator::PooledPhysMem handle); - void pop(); + struct Mapping + { + size_t offset; + PooledPhysMemAllocator::PooledPhysMem handle; + }; + + void push(PooledPhysMemAllocator::PooledPhysMem handle, size_t offset, LocalityDomainId localityDomainId); + void pop(LocalityDomainId localityDomainId); CUdeviceptr mAddr = 0; size_t mVmSize = 0; PooledPhysMemAllocator& mPhysMemAllocator; - std::vector mPhysHandles; + TypedVec> mMappings; + std::set mMappedOffsets; CUmemAccessDesc mAccessDesc{}; }; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/evictionController.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/evictionController.cpp index 7c12e436040d..53175b8e1ca3 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/evictionController.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/evictionController.cpp @@ -76,10 +76,25 @@ LRUEvictionPolicy& PrioritizedEvictionPolicy::getOrCreate(Priority p) return mPolicies[p]; // std::map default-constructs if not present } +void PrioritizedEvictionPolicy::count(Page const& page, SlotCount delta) noexcept +{ + auto const domain = page.localityDomainId(); + if (!domain.has_value()) + { + return; + } + if (*domain >= mDomainCounts.size()) + { + mDomainCounts.resize(*domain + 1, SlotCount{0}); + } + mDomainCounts[*domain] += delta; +} + NodeRef PrioritizedEvictionPolicy::push(SharedPtr page, bool evictFirst) { Priority p = page->priority(); LRUEvictionPolicy& policy = getOrCreate(p); + count(*page, SlotCount{1}); return policy.push(std::move(page), evictFirst); } @@ -93,6 +108,7 @@ SharedPtr PrioritizedEvictionPolicy::pop() { mPolicies.erase(it); } + count(*page, SlotCount{-1}); return page; } @@ -107,6 +123,7 @@ SharedPtr PrioritizedEvictionPolicy::remove(NodeRef node) { mPolicies.erase(it); } + count(*page, SlotCount{-1}); return page; } @@ -120,6 +137,15 @@ SlotCount PrioritizedEvictionPolicy::size() const noexcept return total; } +SlotCount PrioritizedEvictionPolicy::size(std::optional localityDomainId) const noexcept +{ + if (!localityDomainId.has_value()) + { + return size(); + } + return *localityDomainId < mDomainCounts.size() ? mDomainCounts[*localityDomainId] : SlotCount{0}; +} + // allPages() removed — use begin()/end() iterator instead. // --------------------------------------------------------------------------- @@ -169,7 +195,7 @@ void PerLevelEvictionController::scheduleForEviction(Page& page, bool evictFirst } TypedVec>> PerLevelEvictionController::evict( - TypedVec const& minNumPages) + TypedVec const& minNumPages, std::optional localityDomainId) { TLLM_CHECK_DEBUG(minNumPages.size() == numPoolGroups()); @@ -185,16 +211,37 @@ TypedVec>> PerLevelEvictionControlle { throw LogicError("PerLevelEvictionController::evict: page count must be non-negative"); } - SlotCount const available = policy.size() + slotCountValueFromSize(ret[pgIdx].size()); - if (available < count) + if (policy.size(localityDomainId) < count) { throw OutOfPagesError("Not enough pages to evict in group " + std::to_string(pgIdx.value())); } + if (localityDomainId.has_value()) + { + // One pass over the queue: pages of other domains are skipped once, and the + // generator is finished before any removal can invalidate it. + std::vector> candidates; + auto gen = policy.pageGenerator(); + while (slotCountValueFromSize(candidates.size()) < count) + { + auto const* candidate = gen(); + TLLM_CHECK_DEBUG(candidate != nullptr); + if ((*candidate)->localityDomainId() == localityDomainId) + { + candidates.push_back(*candidate); + } + } + for (auto& page : candidates) + { + remove(page->nodeRef.value()); + ret[pgIdx].push_back(std::move(page)); + } + continue; + } while (slotCountValueFromSize(ret[pgIdx].size()) < count) { auto page = policy.pop(); page->nodeRef = std::nullopt; - ret[pgIdx].push_back(page); + ret[pgIdx].push_back(std::move(page)); // @TODO: evict dependencies (like Python _evict_dependencies) } } @@ -235,9 +282,10 @@ void PerLevelEvictionController::remove(NodeRef node) page->nodeRef = std::nullopt; } -SlotCount PerLevelEvictionController::numEvictablePages(PoolGroupIndex pgIdx) const +SlotCount PerLevelEvictionController::numEvictablePages( + PoolGroupIndex pgIdx, std::optional localityDomainId) const { - return mPolicies.at(pgIdx).size(); + return mPolicies.at(pgIdx).size(localityDomainId); } } // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/evictionController.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/evictionController.h index b7ff84816d7f..e9cc83339f0f 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/evictionController.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/evictionController.h @@ -104,6 +104,8 @@ class PrioritizedEvictionPolicy SharedPtr remove(NodeRef node); SlotCount size() const noexcept; + // Pages whose slot lives in the given locality domain; the total when nullopt. + SlotCount size(std::optional localityDomainId) const noexcept; bool empty() const noexcept { @@ -144,9 +146,12 @@ class PrioritizedEvictionPolicy private: LRUEvictionPolicy& getOrCreate(Priority p); + void count(Page const& page, SlotCount delta) noexcept; // std::map keeps keys sorted. We evict from the lowest-priority key first. std::map mPolicies; + // Indexed by locality domain id so a domain-filtered count never walks the queues. + TypedVec mDomainCounts; }; // --------------------------------------------------------------------------- @@ -165,16 +170,18 @@ class PerLevelEvictionController void scheduleForEviction(Page& page, bool evictFirst = false); - // Evict at least minNumPages[pgIdx] pages per pool group. + // Evict at least minNumPages[pgIdx] pages per pool group. With localityDomainId, only pages + // whose slot lives in that domain are taken, in eviction order. // Returns evicted pages per pool group. // On failure, re-queues any already-evicted pages and throws OutOfPagesError. - TypedVec>> evict( - TypedVec const& minNumPages); + TypedVec>> evict(TypedVec const& minNumPages, + std::optional localityDomainId = std::nullopt); // Remove a page from the queue by its NodeRef. void remove(NodeRef node); - [[nodiscard]] SlotCount numEvictablePages(PoolGroupIndex pgIdx) const; + [[nodiscard]] SlotCount numEvictablePages( + PoolGroupIndex pgIdx, std::optional localityDomainId = std::nullopt) const; [[nodiscard]] PoolGroupIndex numPoolGroups() const noexcept { diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp index ee5ed631b8ec..8299dd3b52a9 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp @@ -610,8 +610,8 @@ bool KvCache::resume(std::optional stream, std::optional isDecod throw std::invalid_argument("OPTIONAL reuse selection is only valid before the first resume"); } - // Check utilization against threshold. - auto const utilizations = mManager->storage().getUtilization(kHotLevel); + // Check utilization against threshold. Allocation is per domain, so the guard is too. + auto const utilizations = mManager->storage().getUtilization(kHotLevel, localityDomainId()); float const utilization = utilizations.empty() ? 0.f : *std::max_element(utilizations.begin(), utilizations.end()); if (utilization > mManager->config().maxUtilForResume) { @@ -736,7 +736,7 @@ bool KvCache::resume(std::optional stream, std::optional isDecod CacheLevel dstLevel) { _recordMigratedSlots(pages, slots, srcLevel, dstLevel); }; DropRecorder const dropRecorder = [this](std::vector> const& pages, CacheLevel cacheLevel) { _recordDroppedPages(pages, cacheLevel); }; - tmpSlots = storageMgr.newGpuSlots(numSlotsNeeded, migrationRecorder, dropRecorder); + tmpSlots = storageMgr.newGpuSlots(numSlotsNeeded, migrationRecorder, dropRecorder, localityDomainId()); } catch (OutOfPagesError const&) { @@ -1024,7 +1024,7 @@ bool KvCache::prefetch(CacheLevel target) for (auto const& page : pagesByDestination[destination]) allPages.at(page->lifeCycle).at(page->cacheLevel).push_back(page); - int64_t const diskBlocksMigrated = storageMgr.prefetch(destination, allPages); + int64_t const diskBlocksMigrated = storageMgr.prefetch(destination, allPages, localityDomainId()); if (diskBlocksMigrated > 0 && _shouldRecordManagerStats()) mManager->recordDiskPrefetchBlocks(diskBlocksMigrated); } @@ -1409,7 +1409,7 @@ CommittedPage* KvCache::_copyPageToTreeBlock( Slot newSlot; try { - auto slots = storageMgr.newSlotsForPoolGroup(lvl, pgIdx, 1); + auto slots = storageMgr.newSlotsForPoolGroup(lvl, pgIdx, 1, {}, {}, localityDomainId()); newSlot = std::move(slots[0]); } catch (OutOfPagesError const&) @@ -1781,7 +1781,8 @@ bool KvCache::resize(std::optional capacity, std::optional historyLeng DropRecorder const dropRecorder = [this](std::vector> const& pages, CacheLevel cacheLevel) { _recordDroppedPages(pages, cacheLevel); }; - newSlots = mManager->storage().newGpuSlots(allocCounts, migrationRecorder, dropRecorder); + newSlots + = mManager->storage().newGpuSlots(allocCounts, migrationRecorder, dropRecorder, localityDomainId()); } catch (OutOfPagesError const&) { @@ -2122,7 +2123,7 @@ void KvCache::_appendBeams(BeamIndex oldBeamWidth, BeamIndex newBeamWidth) CacheLevel dstLevel) { _recordMigratedSlots(pages, slots, srcLevel, dstLevel); }; DropRecorder const dropRecorder = [this](std::vector> const& pages, CacheLevel cacheLevel) { _recordDroppedPages(pages, cacheLevel); }; - auto newSlots = storage.newGpuSlots(slotCounts, migrationRecorder, dropRecorder); + auto newSlots = storage.newGpuSlots(slotCounts, migrationRecorder, dropRecorder, localityDomainId()); // Page-index rows must exist before locking the copied pages because lock() // immediately publishes the page index through updateBasePageIndex(). diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h index 612f4bca9bd3..afbd3eb98d3e 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h @@ -462,6 +462,11 @@ class KvCache : public std::enable_shared_from_this return mReuseScope; } + std::optional localityDomainId() const noexcept + { + return mReuseScope.localityDomainId; + } + // Plan dropping SWA blocks needed only by the next conversation turn. // // The plan covers committed pages in each SWA life cycle's current attention diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp index 1ec31e1bc0ae..64bd2088ef9d 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp @@ -192,6 +192,20 @@ std::shared_ptr KvCacheManager::createKvCache(ReuseScope reuseScope, To { KVCM2_API_GUARD(); auto const apiLock = lockExclusive(); + if (numLocalityDomains() > 1) + { + TLLM_CHECK_WITH_INFO(reuseScope.localityDomainId.has_value(), + "locality_domain_id is required when the KV cache manager has %d locality domains", numLocalityDomains()); + LocalityDomainId const localityDomainId = *reuseScope.localityDomainId; + TLLM_CHECK_WITH_INFO(localityDomainId >= 0 && localityDomainId < numLocalityDomains(), + "locality_domain_id %d out of range for %d locality domain(s)", localityDomainId.value(), + numLocalityDomains()); + } + else + { + TLLM_CHECK_WITH_INFO( + !reuseScope.localityDomainId.has_value(), "locality_domain_id is set but locality domains are off"); + } if (!priorityCb) { priorityCb = [](BlockOrdinal, LifeCycleId) { return kPriorityDefault; }; @@ -310,7 +324,7 @@ size_t KvCacheManager::getPageIndexUpperBound(LayerId layerId, DataRole role) co auto const& attr = mStorage->getBufferAttr(layerId, role); LifeCycleId lc = attr.lifeCycleId; PoolGroupIndex pg = mStorage->getPoolGroupIndex(lc); - SlotCount const numSlots = mStorage->numSlots(pg, kHotLevel); + SlotCount const numSlots = mStorage->slotIndexUpperBound(pg, kHotLevel); auto slotSizes = mStorage->slotSize(pg); size_t slotSize = slotSizes.at(attr.poolIndex); return (exactDiv(slotSize, attr.size) * slotCountToSizeT(numSlots) - exactDiv(attr.offset, attr.size)) @@ -436,6 +450,29 @@ int KvCacheManager::tokensPerBlock() const noexcept return mRadixTree->tokensPerBlock(); } +int KvCacheManager::numLocalityDomains() const noexcept +{ + return mStorage->numLocalityDomains(); +} + +std::vector> KvCacheManager::getMappedBufferViews(LayerId layerId, DataRole role) const +{ + auto const apiLock = lockShared(); + auto const& attr = mStorage->getBufferAttr(layerId, role); + PoolGroupIndex const pg = mStorage->getPoolGroupIndex(attr.lifeCycleId); + size_t const slotSize = mStorage->slotSize(pg).at(attr.poolIndex); + MemAddress const base = mStorage->getMemPoolBaseAddress(pg, attr.poolIndex) + attr.offset; + std::vector> views; + for (auto const& [begin, end] : mStorage->slotRanges(pg)) + { + size_t const numPages + = (exactDiv(slotSize, attr.size) * slotCountToSizeT(end - begin) - exactDiv(attr.offset, attr.size)) + * static_cast(attr.expansion); + views.emplace_back(MemAddress(base + slotSize * slotCountToSizeT(begin)), numPages); + } + return views; +} + bool KvCacheManager::enablePartialMatch() const noexcept { return mConfig.enablePartialReuse; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h index 98cff3e0d1cf..ae847198fa62 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h @@ -195,6 +195,10 @@ class KvCacheManager : public std::enable_shared_from_this // ---- Query / info ------------------------------------------------------ int tokensPerBlock() const noexcept; + int numLocalityDomains() const noexcept; + // Contiguous mapped extents of a buffer as (address, page count) pairs. A single pair unless + // the GPU tier is localized, whose pools are mapped in interleaved per-domain chunks. + std::vector> getMappedBufferViews(LayerId layerId, DataRole role) const; bool enablePartialMatch() const noexcept; // Partial commit is independent of partial matching so beam search can diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/localityDomain.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/localityDomain.cpp new file mode 100644 index 000000000000..85c8420173ab --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/localityDomain.cpp @@ -0,0 +1,70 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "kv_cache_manager_v2/localityDomain.h" +#include "kv_cache_manager_v2/exceptions.h" + +#include "tensorrt_llm/runtime/locality_domain/localityDomainResourceConfig.h" +#include "tensorrt_llm/runtime/locality_domain/locality_domain_utils.h" + +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +LocalizationMode localizationMode() +{ + char const* value = std::getenv("TRT_LLM_MOCK_LOCALIZATION_SUPPORT"); + if (value == nullptr) + { + return LocalizationMode::AUTO; + } + std::string_view const mode{value}; + if (mode == "0") + { + return LocalizationMode::OFF; + } + if (mode == "1") + { + return LocalizationMode::MOCK; + } + throw std::invalid_argument("TRT_LLM_MOCK_LOCALIZATION_SUPPORT must be '0' or '1' when set"); +} + +int resolveNumLocalityDomains(bool enableLocalityDomains) +{ + if (!enableLocalityDomains) + { + return 1; + } + switch (localizationMode()) + { + case LocalizationMode::OFF: return 1; + case LocalizationMode::MOCK: return locality_domain::detail::kLocalityDomainCount; + case LocalizationMode::AUTO: + { + CUdevice device{}; + cuCheck(cuCtxGetDevice(&device)); + return locality_domain::deviceSupportsLocalization(device) ? locality_domain::detail::kLocalityDomainCount : 1; + } + } + return 1; +} + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/localityDomain.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/localityDomain.h new file mode 100644 index 000000000000..ac3166aea56c --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/localityDomain.h @@ -0,0 +1,38 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +enum class LocalizationMode +{ + AUTO, + OFF, + MOCK, +}; + +// TRT_LLM_MOCK_LOCALIZATION_SUPPORT: unset probes the device, "0" disables locality domains, +// "1" reports the fixed Rubin domain count on any device while backing every domain with +// ordinary device memory. +LocalizationMode localizationMode(); + +// 1 when locality domains are disabled or the current device has none, otherwise the domain count. +int resolveNumLocalityDomains(bool enableLocalityDomains); + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp index d13e1af69ee7..568b35ec43b3 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp @@ -59,6 +59,7 @@ Page::~Page() { Slot s; s.setSlotId(slotId()); + s.setLocalityDomainId(localityDomainId()); s.readyEvent = std::move(readyEvent); resetSlot(); manager->releaseSlot(lifeCycle, cacheLevel, std::move(s)); @@ -266,6 +267,7 @@ SharedPtr UncommittedPage::convertToCommitted( } // Move slot id to the committed page; invalidate our slot. committed->setSlotId(slotId()); // asserts valid + committed->setLocalityDomainId(localityDomainId()); committed->readyEvent = std::move(readyEvent); resetSlot(); readyEvent = CachedCudaEvent::makeNull(); @@ -671,12 +673,13 @@ std::vector batchedLockPages(KvCache& kvCache, std::vectorcacheLevel != level) ++requirements[storeMgr->getPoolGroupIndex(level, page->lifeCycle)]; - storeMgr->prepareFreeSlots(level, requirements, migrationRecorder, dropRecorder); + storeMgr->prepareFreeSlots(level, requirements, migrationRecorder, dropRecorder, kvCache.localityDomainId()); if (level == kHotLevel && !lockedHostPages.empty()) { - storeMgr->promoteSparsePages(kvCache.cudaStream(), lockedHostPages, migrationRecorder, dropRecorder); + storeMgr->promoteSparsePages( + kvCache.cudaStream(), lockedHostPages, migrationRecorder, dropRecorder, kvCache.localityDomainId()); } - storeMgr->batchedMigrate(level, pages, migrationRecorder); + storeMgr->batchedMigrate(level, pages, migrationRecorder, kvCache.localityDomainId()); } // Wait for all ready events on KvCache's stream (deduplicated). diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.cpp index 0e23097688fc..b87bea96a7bb 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.cpp @@ -51,16 +51,53 @@ size_t tierAllocationGranularity(CacheTier tier, size_t quota) } } +// --------------------------------------------------------------------------- +// SlotIdMap +// --------------------------------------------------------------------------- + +void SlotIdMap::append(size_t chunk) +{ + TLLM_CHECK_DEBUG(mChunkIndices.count(chunk) == 0); + mChunkIndices[chunk] = mChunks.size(); + mChunks.push_back(chunk); +} + +size_t SlotIdMap::pop() +{ + size_t const chunk = mChunks.back(); + mChunks.pop_back(); + mChunkIndices.erase(chunk); + return chunk; +} + +std::optional SlotIdMap::localIndex(SlotId slotId) const noexcept +{ + if (slotId < SlotId{0}) + { + return std::nullopt; + } + auto const chunk = static_cast(slotId.value() / mChunkSize); + auto const it = mChunkIndices.find(chunk); + if (it == mChunkIndices.end()) + { + return std::nullopt; + } + return static_cast(it->second) * mChunkSize + slotId.value() % mChunkSize; +} + // --------------------------------------------------------------------------- // SlotAllocator // --------------------------------------------------------------------------- -SlotAllocator::SlotAllocator(SlotCount capacity) +SlotAllocator::SlotAllocator( + SlotCount capacity, SlotIdMap const* slotMap, std::optional localityDomainId) : mCapacity(capacity) , mTargetCapacity(capacity) , mNumActiveSlots(0) , mNumReadyRecycledSlots(0) , mOccupiedMask(slotCountToSizeT(capacity)) + , mSlotMap(slotMap) + , mLocalityDomainId(localityDomainId) { } @@ -123,7 +160,8 @@ Slot SlotAllocator::allocate() } else if (mNumActiveSlots < std::min(mCapacity, mTargetCapacity)) { - slot.setSlotId(SlotId{mNumActiveSlots++}); + slot.setSlotId(slotId(mNumActiveSlots++)); + slot.setLocalityDomainId(mLocalityDomainId); } else { @@ -131,7 +169,7 @@ Slot SlotAllocator::allocate() mRecycledSlots.pop_front(); TLLM_CHECK_DEBUG_WITH_INFO(slot.hasValidSlot(), "non-ready recycled slot has no valid id"); } - mOccupiedMask.set(toSizeT(slot.slotId())); + mOccupiedMask.set(slotCountToSizeT(*localIndex(slot.slotId()))); return slot; } @@ -160,16 +198,17 @@ void SlotAllocator::release(Slot slot) { throw LogicError("SlotAllocator::release: slot has no valid id"); } - SlotId const slotId = slot.slotId(); - if (!isValidSlotId(slotId) || !mOccupiedMask.get(toSizeT(slotId))) + auto const local = localIndex(slot.slotId()); + if (slot.localityDomainId() != mLocalityDomainId || !local.has_value() || *local < 0 || *local >= mCapacity + || !mOccupiedMask.get(slotCountToSizeT(*local))) { // The id is not one this allocator has outstanding, so drop it rather than let the // rejected slot look like an unreleased one. slot.resetSlot(); throw LogicError("SlotAllocator::release: slot is not occupied"); } - mOccupiedMask.clear(toSizeT(slotId)); - if (slotId < mTargetCapacity) + mOccupiedMask.clear(slotCountToSizeT(*local)); + if (*local < mTargetCapacity) { mRecycledSlots.push_back(std::move(slot)); } @@ -203,7 +242,7 @@ void SlotAllocator::prepareForShrink(SlotCount newNumSlots) SlotCount idx = 0; for (auto& s : mRecycledSlots) { - if (s.slotId() < newNumSlots) + if (*localIndex(s.slotId()) < newNumSlots) { if (idx < oldNumReady) ++newNumReady; @@ -263,7 +302,7 @@ std::vector SlotAllocator::getSlotsBlockingShrink() const for (SlotCount id = mTargetCapacity; id < mCapacity; ++id) { if (mOccupiedMask.get(slotCountToSizeT(id))) - result.push_back(SlotId{id}); + result.push_back(slotId(id)); } return result; } @@ -289,8 +328,8 @@ bool SlotAllocator::check() const noexcept { if (!slot.hasValidSlot()) return false; - SlotCount id = slot.slotId().value(); - if (id < mTargetCapacity || id >= mCapacity) + auto const local = localIndex(slot.slotId()); + if (!local.has_value() || *local < mTargetCapacity || *local >= mCapacity) return false; } SlotCount const accountedSlots = slotCountValueFromSize(mRecycledSlots.size()) @@ -319,13 +358,20 @@ void SlotAllocator::scrubEvents() // GpuSlotPool // --------------------------------------------------------------------------- -GpuSlotPool::GpuSlotPool(size_t slotSize, size_t vmSize, PooledPhysMemAllocator& physMemAllocator, SlotCount numSlots) +GpuSlotPool::GpuSlotPool(size_t slotSize, size_t vmSize, PooledPhysMemAllocator& physMemAllocator, SlotCount numSlots, + TypedVec const* slotMaps) : SlotPoolBase(slotSize) , mVirtMem(vmSize, physMemAllocator) + , mSlotMaps(slotMaps) { TLLM_CHECK_DEBUG_WITH_INFO( vmSize % physMemAllocator.physMemSize() == 0, "vm_size must be aligned to phys_mem_size"); - resize(numSlots); + TLLM_CHECK_DEBUG( + mSlotMaps == nullptr || mSlotMaps->size() == LocalityDomainId{physMemAllocator.numLocalityDomains()}); + if (numSlots != 0) + { + resize(numSlots); + } } size_t GpuSlotPool::computeNumPhysMem(size_t slotSize, SlotCount numSlots, size_t physMemSize) noexcept @@ -340,7 +386,17 @@ SlotCount GpuSlotPool::computeNumSlots(size_t slotSize, size_t numPhysMem, size_ SlotCount GpuSlotPool::numSlots() const noexcept { - return computeNumSlots(mSlotSize, mVirtMem.numPhysMem(), mVirtMem.physMemSize()); + SlotCount total = 0; + for (LocalityDomainId localityDomainId{0}; localityDomainId < numLocalityDomains(); ++localityDomainId) + { + total += numSlots(localityDomainId); + } + return total; +} + +SlotCount GpuSlotPool::numSlots(LocalityDomainId localityDomainId) const noexcept +{ + return computeNumSlots(mSlotSize, mVirtMem.numPhysMem(localityDomainId), mVirtMem.physMemSize()); } void GpuSlotPool::destroy() @@ -350,9 +406,55 @@ void GpuSlotPool::destroy() void GpuSlotPool::resize(SlotCount newNumSlots) { - size_t physSize = mVirtMem.physMemSize(); - size_t newPhysMem = computeNumPhysMem(mSlotSize, newNumSlots, physSize); - mVirtMem.realloc(physSize * newPhysMem); + resize(newNumSlots, LocalityDomainId{0}); +} + +void GpuSlotPool::resize(SlotCount newNumSlots, LocalityDomainId localityDomainId) +{ + size_t const physSize = mVirtMem.physMemSize(); + size_t const required = computeNumPhysMem(mSlotSize, newNumSlots, physSize); + size_t const old = mVirtMem.numPhysMem(localityDomainId); + if (required < old) + { + mVirtMem.shrink(old - required, localityDomainId); + } + else if (required > old) + { + try + { + if (mSlotMaps == nullptr) + { + mVirtMem.extend(required - old, localityDomainId, old * physSize); + } + else + { + // Walk the domain's chunks in order; a chunk's bytes are a multiple of physSize in + // every pool, so each piece maps whole physical chunks. + auto const& slotMap = (*mSlotMaps)[localityDomainId]; + size_t const chunkBytes = slotCountToSizeT(slotMap.chunkSize()) * mSlotSize; + size_t mapped = old * physSize; + size_t const end = required * physSize; + while (mapped < end) + { + size_t const chunkIndex = mapped / chunkBytes; + size_t const offsetInChunk = mapped % chunkBytes; + size_t const size = std::min(end - mapped, chunkBytes - offsetInChunk); + size_t const offset = slotMap.chunks().at(chunkIndex) * chunkBytes + offsetInChunk; + mVirtMem.extend(size / physSize, localityDomainId, offset); + mapped += size; + } + } + } + catch (...) + { + size_t const added = mVirtMem.numPhysMem(localityDomainId) - old; + if (added != 0) + { + mVirtMem.shrink(added, localityDomainId); + } + throw; + } + } } SlotCount GpuSlotPool::extendByOnePhysMem() @@ -363,7 +465,11 @@ SlotCount GpuSlotPool::extendByOnePhysMem() Address GpuSlotPool::slotAddress(SlotId slot) const { - TLLM_CHECK_DEBUG_WITH_INFO(isValidSlotId(slot), "GpuSlotPool::slotAddress: slot index out of bounds"); + // Slot ids may lie beyond numSlots() when domains interleave, so bound by the reservation. + TLLM_CHECK_DEBUG_WITH_INFO(mSlotMaps == nullptr + ? isValidSlotId(slot) + : slot >= SlotId{0} && mSlotSize * (toSizeT(slot) + 1) <= mVirtMem.virtualBytes(), + "GpuSlotPool::slotAddress: slot index out of bounds"); return MemAddress(mVirtMem.address() + mSlotSize * toSizeT(slot)); } @@ -498,11 +604,11 @@ Address DiskSlotPool::slotAddress(SlotId slot) const // --------------------------------------------------------------------------- PoolGroupBase::PoolGroupBase(SlotCount numSlots) - : mSlotAllocator(numSlots) { + mSlotAllocators.push_back(std::make_unique(numSlots)); } -SlotCount PoolGroupBase::getNumSlotsFromPools() const noexcept +SlotCount PoolGroupBase::getNumSlotsFromPools(LocalityDomainId /*localityDomainId*/) const noexcept { if (mPools.empty()) return 0; @@ -519,59 +625,74 @@ PoolGroupBase::~PoolGroupBase() KVCM2_POISON_ON_EXCEPT([this]() { destroy(); }); } -SlotCount PoolGroupBase::numSlots() const noexcept +SlotCount PoolGroupBase::numSlots(LocalityDomainId localityDomainId) const noexcept { - SlotCount n = mSlotAllocator.numSlots(); + SlotCount n = mSlotAllocators[localityDomainId]->numSlots(); if (TLLM_UNLIKELY(gDebug)) { // Mirrors Python PoolGroupBase.num_slots: assert num_slots <= self._get_num_slots_from_pools() - [[maybe_unused]] SlotCount poolSlots = getNumSlotsFromPools(); + [[maybe_unused]] SlotCount poolSlots = getNumSlotsFromPools(localityDomainId); TLLM_CHECK_WITH_INFO(n <= poolSlots, "SlotAllocator capacity exceeds pool capacity"); } return n; } -Slot PoolGroupBase::allocate() +Slot PoolGroupBase::allocate(LocalityDomainId localityDomainId) { - return mSlotAllocator.allocate(); + return mSlotAllocators.at(localityDomainId)->allocate(); } -std::vector PoolGroupBase::allocateMultiple(SlotCount numSlots) +std::vector PoolGroupBase::allocateMultiple(SlotCount numSlots, LocalityDomainId localityDomainId) { - return mSlotAllocator.allocateMultiple(numSlots); + return mSlotAllocators.at(localityDomainId)->allocateMultiple(numSlots); } void PoolGroupBase::release(Slot slot) { - mSlotAllocator.release(std::move(slot)); + LocalityDomainId const owner = slot.localityDomainId().value_or(LocalityDomainId{0}); + mSlotAllocators.at(owner)->release(std::move(slot)); } void PoolGroupBase::destroy() { if (mDestroyed) return; - if (mSlotAllocator.numSlots() != 0) + for (auto& allocator : mSlotAllocators) { - mSlotAllocator.synchronize(); - mSlotAllocator.prepareForShrink(0); - mSlotAllocator.finishShrink(); + if (allocator->numSlots() != 0) + { + allocator->synchronize(); + allocator->prepareForShrink(0); + allocator->finishShrink(); + } } for (auto& p : mPools) p->destroy(); mDestroyed = true; } -void PoolGroupBase::resizePools(std::optional newNumSlots) +void PoolGroupBase::resizePools(std::optional newNumSlots, LocalityDomainId localityDomainId) { - SlotCount n = newNumSlots.value_or(mSlotAllocator.numSlots()); + TLLM_CHECK_DEBUG(localityDomainId == LocalityDomainId{0}); + SlotCount n = newNumSlots.value_or(mSlotAllocators[LocalityDomainId{0}]->numSlots()); for (auto& p : mPools) p->resize(n); // Mirrors Python PoolGroupBase.resize_pools: assert NDEBUG or self._check(True) // After resize, allocator capacity must not exceed pool capacity (allow mismatch). - TLLM_CHECK_DEBUG_WITH_INFO(mSlotAllocator.numSlots() <= getNumSlotsFromPools(), + TLLM_CHECK_DEBUG_WITH_INFO(mSlotAllocators[LocalityDomainId{0}]->numSlots() <= getNumSlotsFromPools(), "After resizePools: allocator capacity exceeds pool capacity"); } +std::vector> PoolGroupBase::slotRanges() const +{ + SlotCount const count = mSlotAllocators[LocalityDomainId{0}]->numSlots(); + if (count == 0) + { + return {}; + } + return {{SlotCount{0}, count}}; +} + TypedVec PoolGroupBase::slotAddress(SlotId slotId) const { TypedVec addrs; @@ -594,10 +715,24 @@ TypedVec PoolGroupBase::slotSize() const // GpuPoolGroup // --------------------------------------------------------------------------- -GpuPoolGroup::GpuPoolGroup( - SlotCount numSlots, TypedVec const& slotSizeList, PooledPhysMemAllocator& physMemAllocator) - : PoolGroupBase(numSlots) +size_t localityDomainChunkSlots(TypedVec const& slotSizeList, size_t physMemSize) { + // Includes the pools holding smaller quantization-scale buffers. + size_t chunkSlots = 1; + for (size_t sz : slotSizeList) + { + chunkSlots = std::lcm(chunkSlots, physMemSize / std::gcd(physMemSize, sz)); + } + return chunkSlots; +} + +GpuPoolGroup::GpuPoolGroup(TypedVec const& numSlots, + TypedVec const& slotSizeList, PooledPhysMemAllocator& physMemAllocator) + : PoolGroupBase(physMemAllocator.numLocalityDomains() == 1 ? numSlots.at(LocalityDomainId{0}) : SlotCount{0}) +{ + int const numLocalityDomains = physMemAllocator.numLocalityDomains(); + TLLM_CHECK_WITH_INFO(numSlots.size().value() == numLocalityDomains, + "GpuPoolGroup: one initial slot count per locality domain is required"); size_t physMemSize = physMemAllocator.physMemSize(); // Query total GPU memory to size virtual address space (mirrors Python GpuPoolGroup). size_t totalGpuMem = 0; @@ -609,15 +744,196 @@ GpuPoolGroup::GpuPoolGroup( // @TODO: We should replace maxSlotSize with sum. This should also be updated in Python. Will do it later. TLLM_CHECK_WITH_INFO(!slotSizeList.empty(), "GpuPoolGroup: slotSizeList must not be empty"); size_t maxSlotSize = *std::max_element(slotSizeList.begin(), slotSizeList.end()); + + if (numLocalityDomains == 1) + { + for (size_t sz : slotSizeList) + { + // VA proportional to GPU memory, scaled by slot size ratio (mirrors Python). + // Compute ratio as double first to avoid size_t overflow. + double sizeRatio = static_cast(sz) / static_cast(maxSlotSize); + size_t vmSize = roundDown(static_cast(static_cast(totalGpuMem) * sizeRatio), physMemSize); + vmSize = std::max(vmSize, roundUp(slotCountToSizeT(numSlots[LocalityDomainId{0}]) * sz, physMemSize)); + mPools.push_back( + std::make_unique(sz, vmSize, physMemAllocator, numSlots[LocalityDomainId{0}])); + } + return; + } + + size_t const chunkSize = localityDomainChunkSlots(slotSizeList, physMemSize); + mChunkSize = static_cast(chunkSize); + size_t const maxSlots = totalGpuMem / maxSlotSize; + // At most one partially used chunk per domain needs extra VA padding. + mMaxChunks = divUp(maxSlots, chunkSize) + static_cast(numLocalityDomains); + + mSlotMaps.reserve(LocalityDomainId{numLocalityDomains}); + mSlotAllocators.clear(); + for (LocalityDomainId localityDomainId{0}; localityDomainId < numLocalityDomains; ++localityDomainId) + { + mSlotMaps.emplace_back(mChunkSize); + } + for (LocalityDomainId localityDomainId{0}; localityDomainId < numLocalityDomains; ++localityDomainId) + { + mSlotAllocators.push_back( + std::make_unique(SlotCount{0}, &mSlotMaps[localityDomainId], localityDomainId)); + } for (size_t sz : slotSizeList) { - // VA proportional to GPU memory, scaled by slot size ratio (mirrors Python). - // Compute ratio as double first to avoid size_t overflow. - double sizeRatio = static_cast(sz) / static_cast(maxSlotSize); - size_t vmSize = roundDown(static_cast(static_cast(totalGpuMem) * sizeRatio), physMemSize); - vmSize = std::max(vmSize, roundUp(slotCountToSizeT(numSlots) * sz, physMemSize)); - mPools.push_back(std::make_unique(sz, vmSize, physMemAllocator, numSlots)); + mPools.push_back( + std::make_unique(sz, mMaxChunks * chunkSize * sz, physMemAllocator, SlotCount{0}, &mSlotMaps)); + } + for (LocalityDomainId localityDomainId{0}; localityDomainId < numLocalityDomains; ++localityDomainId) + { + SlotCount const count = numSlots[localityDomainId]; + resizePools(count, localityDomainId); + if (count != 0) + { + mSlotAllocators[localityDomainId]->expand(count); + } + } +} + +GpuPoolGroup::~GpuPoolGroup() +{ + KVCM2_POISON_ON_EXCEPT([this]() { destroy(); }); +} + +size_t GpuPoolGroup::roundedNumBytes(size_t granularity) const noexcept +{ + size_t total = 0; + for (auto const& pool : mPools) + { + auto const& gpuPool = static_cast(*pool); + for (LocalityDomainId localityDomainId{0}; localityDomainId < gpuPool.numLocalityDomains(); ++localityDomainId) + { + total += roundUp(gpuPool.slotSize() * slotCountToSizeT(gpuPool.numSlots(localityDomainId)), granularity); + } + } + return total; +} + +SlotCount GpuPoolGroup::getNumSlotsFromPools(LocalityDomainId localityDomainId) const noexcept +{ + if (mPools.empty()) + return 0; + SlotCount minSlots = std::numeric_limits::max(); + for (auto const& pool : mPools) + { + minSlots = std::min(minSlots, static_cast(*pool).numSlots(localityDomainId)); + } + return minSlots; +} + +void GpuPoolGroup::resizeSlotMap(SlotCount numSlots, LocalityDomainId localityDomainId) +{ + if (mSlotMaps.empty()) + { + return; + } + auto& slotMap = mSlotMaps[localityDomainId]; + size_t const numChunks = divUp(slotCountToSizeT(numSlots), slotCountToSizeT(mChunkSize)); + if (numChunks > slotMap.chunks().size() + mFreeChunks.size() + mMaxChunks - mNumChunks) + { + throw OutOfPagesError("GpuPoolGroup: slot ranges exceed the virtual memory reservation"); + } + while (slotMap.chunks().size() < numChunks) + { + size_t chunk; + if (!mFreeChunks.empty()) + { + chunk = mFreeChunks.top(); + mFreeChunks.pop(); + } + else + { + chunk = mNumChunks++; + } + slotMap.append(chunk); + } + while (slotMap.chunks().size() > numChunks) + { + mFreeChunks.push(slotMap.pop()); + } +} + +// Resize backing transactionally without changing issued ids or the allocator capacity. +void GpuPoolGroup::resizePools(std::optional newNumSlots, LocalityDomainId localityDomainId) +{ + if (mSlotMaps.empty()) + { + PoolGroupBase::resizePools(newNumSlots, localityDomainId); + return; + } + SlotCount const target = newNumSlots.value_or(mSlotAllocators.at(localityDomainId)->numSlots()); + std::vector oldCounts; + oldCounts.reserve(mPools.stdSize()); + SlotCount oldNumSlots = 0; + for (auto const& pool : mPools) + { + SlotCount const count = static_cast(*pool).numSlots(localityDomainId); + oldCounts.push_back(count); + oldNumSlots = std::max(oldNumSlots, count); + } + if (target > oldNumSlots) + { + resizeSlotMap(target, localityDomainId); + } + try + { + for (auto& pool : mPools) + { + static_cast(*pool).resize(target, localityDomainId); + } + } + catch (...) + { + for (PoolIndex poolIdx{0}; poolIdx < mPools.size(); ++poolIdx) + { + static_cast(*mPools[poolIdx]).resize(oldCounts[toSizeT(poolIdx)], localityDomainId); + } + resizeSlotMap(oldNumSlots, localityDomainId); + throw; + } + resizeSlotMap(target, localityDomainId); + TLLM_CHECK_DEBUG_WITH_INFO(mSlotAllocators[localityDomainId]->numSlots() <= getNumSlotsFromPools(localityDomainId), + "After resizePools: allocator capacity exceeds pool capacity"); +} + +std::vector> GpuPoolGroup::slotRanges() const +{ + if (mSlotMaps.empty()) + { + return PoolGroupBase::slotRanges(); + } + std::vector> ranges; + for (LocalityDomainId localityDomainId{0}; localityDomainId < numLocalityDomains(); ++localityDomainId) + { + SlotCount const count = mSlotAllocators[localityDomainId]->numSlots(); + auto const& chunks = mSlotMaps[localityDomainId].chunks(); + for (size_t i = 0; i < chunks.size(); ++i) + { + SlotCount const size = std::min(mChunkSize, count - static_cast(i) * mChunkSize); + if (size > 0) + { + SlotCount const start = static_cast(chunks[i]) * mChunkSize; + ranges.emplace_back(start, start + size); + } + } + } + std::sort(ranges.begin(), ranges.end()); + std::vector> merged; + for (auto const& [start, end] : ranges) + { + if (!merged.empty() && merged.back().second == start) + { + merged.back().second = end; + } + else + { + merged.emplace_back(start, end); + } } + return merged; } // --------------------------------------------------------------------------- @@ -651,9 +967,9 @@ void HostPoolGroup::waitSlotUsable(SlotId slot) const // A failed fill surfaces out of waitSlotUsable, after the slots are already // taken, so they are returned to the allocator before the error propagates. -Slot HostPoolGroup::allocate() +Slot HostPoolGroup::allocate(LocalityDomainId localityDomainId) { - Slot slot = PoolGroupBase::allocate(); + Slot slot = PoolGroupBase::allocate(localityDomainId); try { waitSlotUsable(slot.slotId()); @@ -666,9 +982,9 @@ Slot HostPoolGroup::allocate() return slot; } -std::vector HostPoolGroup::allocateMultiple(SlotCount numSlots) +std::vector HostPoolGroup::allocateMultiple(SlotCount numSlots, LocalityDomainId localityDomainId) { - std::vector slots = PoolGroupBase::allocateMultiple(numSlots); + std::vector slots = PoolGroupBase::allocateMultiple(numSlots, localityDomainId); if (slots.empty()) { return slots; @@ -712,9 +1028,10 @@ DiskPoolGroup::DiskPoolGroup( // CacheLevelStorage // --------------------------------------------------------------------------- -std::vector CacheLevelStorage::allocateMultiple(PoolGroupIndex pgIdx, SlotCount numSlots) +std::vector CacheLevelStorage::allocateMultiple( + PoolGroupIndex pgIdx, SlotCount numSlots, LocalityDomainId localityDomainId) { - return mPoolGroups.at(pgIdx)->allocateMultiple(numSlots); + return mPoolGroups.at(pgIdx)->allocateMultiple(numSlots, localityDomainId); } void CacheLevelStorage::release(PoolGroupIndex pgIdx, Slot slot) @@ -722,9 +1039,9 @@ void CacheLevelStorage::release(PoolGroupIndex pgIdx, Slot slot) mPoolGroups.at(pgIdx)->release(std::move(slot)); } -SlotCount CacheLevelStorage::numFreeSlots(PoolGroupIndex pgIdx) const +SlotCount CacheLevelStorage::numFreeSlots(PoolGroupIndex pgIdx, LocalityDomainId localityDomainId) const { - return mPoolGroups.at(pgIdx)->numFreeSlots(); + return mPoolGroups.at(pgIdx)->numFreeSlots(localityDomainId); } TypedVec CacheLevelStorage::slotAddress(PoolGroupIndex pgIdx, SlotId slotId) const @@ -742,10 +1059,13 @@ GpuCacheLevelStorage::GpuCacheLevelStorage(TypedVec co { TLLM_CHECK_WITH_INFO(slotCountList.size() == slotDescList.size(), "GpuCacheLevelStorage: slotCountList and slotDescList must have the same length"); + // slotCountList is per domain: StorageManager sizes it from an equal share of the quota. + int const numLocalityDomains = mPhysMemAllocator.numLocalityDomains(); for (PoolGroupIndex pgIdx{0}; pgIdx < slotDescList.size(); ++pgIdx) { mPoolGroups.push_back(std::make_unique( - slotCountList[pgIdx], slotDescList[pgIdx].slotSizeList(), mPhysMemAllocator)); + TypedVec(LocalityDomainId{numLocalityDomains}, slotCountList[pgIdx]), + slotDescList[pgIdx].slotSizeList(), mPhysMemAllocator)); } } diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.h index 5f4c5c8da375..f5109c891b59 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.h @@ -29,6 +29,8 @@ #include #include #include +#include +#include #include #include @@ -61,6 +63,7 @@ struct Slot Slot(Slot&& other) noexcept : readyEvent(std::move(other.readyEvent)) , mSlotId(std::exchange(other.mSlotId, std::nullopt)) + , mLocalityDomainId(std::exchange(other.mLocalityDomainId, std::nullopt)) { } @@ -70,6 +73,7 @@ struct Slot { readyEvent = std::move(other.readyEvent); mSlotId = std::exchange(other.mSlotId, std::nullopt); + mLocalityDomainId = std::exchange(other.mLocalityDomainId, std::nullopt); } return *this; } @@ -93,9 +97,21 @@ struct Slot mSlotId = id; } + // Locality domain whose physical memory backs this slot; nullopt for non-localized storage. + [[nodiscard]] std::optional localityDomainId() const noexcept + { + return mLocalityDomainId; + } + + void setLocalityDomainId(std::optional id) noexcept + { + mLocalityDomainId = id; + } + void resetSlot() noexcept { mSlotId.reset(); + mLocalityDomainId.reset(); } bool queryReady() @@ -121,22 +137,70 @@ struct Slot Slot previous; auto replacementId = std::exchange(replacement.mSlotId, std::nullopt); previous.mSlotId = std::exchange(mSlotId, replacementId); + auto replacementDomain = std::exchange(replacement.mLocalityDomainId, std::nullopt); + previous.mLocalityDomainId = std::exchange(mLocalityDomainId, replacementDomain); previous.readyEvent = std::exchange(readyEvent, std::move(replacement.readyEvent)); return previous; } private: std::optional mSlotId; + std::optional mLocalityDomainId; +}; + +// --------------------------------------------------------------------------- +// SlotIdMap — maps an allocator's local slot indices onto chunks of a slot-id space shared by +// several allocators (one per locality domain), so each hands out disjoint global ids. +// Mirrors _storage/_core.py::_SlotIdMap. +// --------------------------------------------------------------------------- +class SlotIdMap +{ +public: + explicit SlotIdMap(SlotCount chunkSize) + : mChunkSize(chunkSize) + { + } + + [[nodiscard]] SlotCount chunkSize() const noexcept + { + return mChunkSize; + } + + [[nodiscard]] std::vector const& chunks() const noexcept + { + return mChunks; + } + + void append(size_t chunk); + size_t pop(); + + [[nodiscard]] SlotId slotId(SlotCount localIdx) const noexcept + { + return SlotId{static_cast(mChunks[slotCountToSizeT(localIdx / mChunkSize)]) * mChunkSize + + localIdx % mChunkSize}; + } + + // nullopt when slotId lies in a chunk this map does not own. + [[nodiscard]] std::optional localIndex(SlotId slotId) const noexcept; + +private: + SlotCount mChunkSize; + std::vector mChunks; + std::unordered_map mChunkIndices; }; // --------------------------------------------------------------------------- // SlotAllocator — manages a fixed-capacity array of slot ids. // Mirrors _storage/_core.py::SlotAllocator. +// +// Without a SlotIdMap the slot id is the local index. With one, capacity and shrink accounting +// stay local while the ids handed out come from the map's chunks. // --------------------------------------------------------------------------- class SlotAllocator { public: - explicit SlotAllocator(SlotCount capacity); + explicit SlotAllocator(SlotCount capacity, SlotIdMap const* slotMap = nullptr, + std::optional localityDomainId = std::nullopt); ~SlotAllocator(); [[nodiscard]] SlotCount numFreeSlots() const noexcept; @@ -147,11 +211,23 @@ class SlotAllocator return mCapacity; } + [[nodiscard]] std::optional localityDomainId() const noexcept + { + return mLocalityDomainId; + } + + // Local index of a slot this allocator issued; nullopt for a foreign slot id. + [[nodiscard]] std::optional localIndexOf(SlotId slot) const noexcept + { + return localIndex(slot); + } + // SlotId is signed, so a bare `slot < numSlots()` admits negatives. Both ends are checked // here so no caller has to remember the lower one. [[nodiscard]] bool isValidSlotId(SlotId slot) const noexcept { - return slot >= SlotId{0} && slot < numSlots(); + auto const local = localIndex(slot); + return local.has_value() && *local >= 0 && *local < numSlots(); } Slot allocate(); @@ -191,13 +267,25 @@ class SlotAllocator void scrubEvents(); [[nodiscard]] bool check() const noexcept; + [[nodiscard]] std::optional localIndex(SlotId slot) const noexcept + { + return mSlotMap == nullptr ? std::optional{slot.value()} : mSlotMap->localIndex(slot); + } + + [[nodiscard]] SlotId slotId(SlotCount localIdx) const noexcept + { + return mSlotMap == nullptr ? SlotId{localIdx} : mSlotMap->slotId(localIdx); + } + SlotCount mCapacity; SlotCount mTargetCapacity; SlotCount mNumActiveSlots; SlotCount mNumReadyRecycledSlots; std::deque mRecycledSlots; std::vector mOverflowSlots; - DynamicBitset mOccupiedMask; + DynamicBitset mOccupiedMask; // indexed by local index + SlotIdMap const* mSlotMap; + std::optional mLocalityDomainId; }; //! Granularity a disk pool is sized in. Disk has no allocation unit to match, @@ -264,13 +352,30 @@ class SlotPoolBase class GpuSlotPool : public SlotPoolBase { public: - GpuSlotPool(size_t slotSize, size_t vmSize, PooledPhysMemAllocator& physMemAllocator, SlotCount numSlots); + // slotMaps: one per locality domain when the group places each domain's backing at + // chunk-aligned offsets of a shared slot-id space; nullptr for a single contiguous mapping. + GpuSlotPool(size_t slotSize, size_t vmSize, PooledPhysMemAllocator& physMemAllocator, SlotCount numSlots, + TypedVec const* slotMaps = nullptr); + // Slots backed across all locality domains. SlotCount numSlots() const noexcept override; + SlotCount numSlots(LocalityDomainId localityDomainId) const noexcept; void destroy() override; void resize(SlotCount newNumSlots) override; + void resize(SlotCount newNumSlots, LocalityDomainId localityDomainId); Address slotAddress(SlotId slot) const override; + [[nodiscard]] int numLocalityDomains() const noexcept + { + return mVirtMem.numLocalityDomains(); + } + + // Merged mapped byte ranges relative to the pool base. + [[nodiscard]] std::vector> mappedRanges() const + { + return mVirtMem.mappedRanges(); + } + // Extend by exactly one physical memory chunk; returns new numSlots. SlotCount extendByOnePhysMem(); @@ -279,6 +384,7 @@ class GpuSlotPool : public SlotPoolBase private: VirtMem mVirtMem; + TypedVec const* mSlotMaps; }; // --------------------------------------------------------------------------- @@ -354,28 +460,50 @@ class PoolGroupBase return mPools.size(); } - SlotCount numSlots() const noexcept; + int numLocalityDomains() const noexcept + { + return mSlotAllocators.size().value(); + } - SlotCount numFreeSlots() const noexcept + SlotCount numSlots(LocalityDomainId localityDomainId = LocalityDomainId{0}) const noexcept; + + SlotCount numFreeSlots(LocalityDomainId localityDomainId = LocalityDomainId{0}) const noexcept { - return mSlotAllocator.numFreeSlots(); + return mSlotAllocators[localityDomainId]->numFreeSlots(); } - SlotAllocator& slotAllocator() noexcept + SlotAllocator& slotAllocator(LocalityDomainId localityDomainId = LocalityDomainId{0}) noexcept { - return mSlotAllocator; + return *mSlotAllocators[localityDomainId]; } - SlotAllocator const& slotAllocator() const noexcept + SlotAllocator const& slotAllocator(LocalityDomainId localityDomainId = LocalityDomainId{0}) const noexcept { - return mSlotAllocator; + return *mSlotAllocators[localityDomainId]; } - virtual Slot allocate(); - virtual std::vector allocateMultiple(SlotCount numSlots); + virtual Slot allocate(LocalityDomainId localityDomainId = LocalityDomainId{0}); + virtual std::vector allocateMultiple( + SlotCount numSlots, LocalityDomainId localityDomainId = LocalityDomainId{0}); + // Returns the slot to the allocator of the domain that issued it. void release(Slot slot); void destroy(); - void resizePools(std::optional newNumSlots = std::nullopt); + virtual void resizePools( + std::optional newNumSlots = std::nullopt, LocalityDomainId localityDomainId = LocalityDomainId{0}); + + // Usable global slot ranges [begin, end), merged and sorted, excluding unmapped holes. + virtual std::vector> slotRanges() const; + + // One past the highest usable slot id; equals numSlots() for a contiguous group. + SlotCount slotIndexUpperBound() const + { + SlotCount upperBound = 0; + for (auto const& [begin, end] : slotRanges()) + { + upperBound = std::max(upperBound, end); + } + return upperBound; + } // Addresses for all pools at a given slot id. TypedVec slotAddress(SlotId slotId) const; @@ -393,7 +521,7 @@ class PoolGroupBase // Total bytes across all pools, rounding each pool's bytes up to granularity. // Mirrors Python total_quota: sum(round_up(p.num_bytes, granularity) for p in pg._pools). - size_t roundedNumBytes(size_t granularity) const noexcept + virtual size_t roundedNumBytes(size_t granularity) const noexcept { size_t total = 0; for (auto const& pool : mPools) @@ -403,21 +531,56 @@ class PoolGroupBase protected: // Mirrors Python _get_num_slots_from_pools: min(p.num_slots for p in pools). - SlotCount getNumSlotsFromPools() const noexcept; + virtual SlotCount getNumSlotsFromPools(LocalityDomainId localityDomainId = LocalityDomainId{0}) const noexcept; - SlotAllocator mSlotAllocator; + TypedVec> mSlotAllocators; TypedVec> mPools; bool mDestroyed = false; }; +// Slots per locality-domain chunk: the smallest slot count that starts on a physical-page boundary in every pool +// of the group. +[[nodiscard]] size_t localityDomainChunkSlots(TypedVec const& slotSizeList, size_t physMemSize); + // --------------------------------------------------------------------------- // GpuPoolGroup / HostPoolGroup / DiskPoolGroup // --------------------------------------------------------------------------- + +// With several locality domains the pools share one slot-id space: ids are handed out in chunks +// whose byte extent is physically aligned in every pool, each chunk backed by one domain's memory. class GpuPoolGroup : public PoolGroupBase { public: - GpuPoolGroup( - SlotCount numSlots, TypedVec const& slotSizeList, PooledPhysMemAllocator& physMemAllocator); + // numSlots: initial slot count per locality domain (size numLocalityDomains()). + GpuPoolGroup(TypedVec const& numSlots, TypedVec const& slotSizeList, + PooledPhysMemAllocator& physMemAllocator); + // The allocators hold pointers into mSlotMaps, so release them before the maps go away. + ~GpuPoolGroup() override; + + void resizePools(std::optional newNumSlots = std::nullopt, + LocalityDomainId localityDomainId = LocalityDomainId{0}) override; + std::vector> slotRanges() const override; + // Each domain's backing is rounded up on its own, so the sum must be too or a quota + // round trip through totalQuota() loses slots. + size_t roundedNumBytes(size_t granularity) const noexcept override; + + [[nodiscard]] std::vector> mappedRanges(PoolIndex poolIdx) const + { + return static_cast(*mPools.at(poolIdx)).mappedRanges(); + } + +protected: + SlotCount getNumSlotsFromPools(LocalityDomainId localityDomainId = LocalityDomainId{0}) const noexcept override; + +private: + // Grow or shrink a domain's chunk list to cover numSlots. + void resizeSlotMap(SlotCount numSlots, LocalityDomainId localityDomainId); + + TypedVec mSlotMaps; // empty with one locality domain + SlotCount mChunkSize{0}; + size_t mMaxChunks{0}; + size_t mNumChunks{0}; + std::priority_queue, std::greater<>> mFreeChunks; }; //! Host pools are committed progressively on a background thread, so a slot may @@ -439,8 +602,9 @@ class HostPoolGroup : public PoolGroupBase HostPoolGroup(SlotCount numSlots, TypedVec const& slotSizeList, size_t groupVmSize, HostMemBackingOptions const& options = {}); - Slot allocate() override; - std::vector allocateMultiple(SlotCount numSlots) override; + Slot allocate(LocalityDomainId localityDomainId = LocalityDomainId{0}) override; + std::vector allocateMultiple( + SlotCount numSlots, LocalityDomainId localityDomainId = LocalityDomainId{0}) override; private: //! Blocks until every pool backs the byte range ending just past `slot`. @@ -469,11 +633,18 @@ class CacheLevelStorage return mPoolGroups.size(); } - std::vector allocateMultiple(PoolGroupIndex pgIdx, SlotCount numSlots); + std::vector allocateMultiple( + PoolGroupIndex pgIdx, SlotCount numSlots, LocalityDomainId localityDomainId = LocalityDomainId{0}); void release(PoolGroupIndex pgIdx, Slot slot); - SlotCount numFreeSlots(PoolGroupIndex pgIdx) const; + SlotCount numFreeSlots(PoolGroupIndex pgIdx, LocalityDomainId localityDomainId = LocalityDomainId{0}) const; TypedVec slotAddress(PoolGroupIndex pgIdx, SlotId slotId) const; + // 1 unless this is a localized GPU level. + virtual int numLocalityDomains() const noexcept + { + return 1; + } + // Additional accessors used by StorageManager and KvCacheManager. virtual void destroy() { @@ -486,9 +657,19 @@ class CacheLevelStorage return mPoolGroups.at(pgIdx)->numPools(); } - SlotCount numSlots(PoolGroupIndex pgIdx) const + SlotCount numSlots(PoolGroupIndex pgIdx, LocalityDomainId localityDomainId = LocalityDomainId{0}) const + { + return mPoolGroups.at(pgIdx)->numSlots(localityDomainId); + } + + SlotCount slotIndexUpperBound(PoolGroupIndex pgIdx) const { - return mPoolGroups.at(pgIdx)->numSlots(); + return mPoolGroups.at(pgIdx)->slotIndexUpperBound(); + } + + std::vector> slotRanges(PoolGroupIndex pgIdx) const + { + return mPoolGroups.at(pgIdx)->slotRanges(); } TypedVec slotSize(PoolGroupIndex pgIdx) const @@ -520,13 +701,13 @@ class CacheLevelStorage return total; } - // Returns numSlots() per pool group. - TypedVec slotCountList() const + // Returns numSlots(localityDomainId) per pool group. + TypedVec slotCountList(LocalityDomainId localityDomainId = LocalityDomainId{0}) const { TypedVec ret; ret.reserve(mPoolGroups.size()); for (auto const& pg : mPoolGroups) - ret.push_back(pg->numSlots()); + ret.push_back(pg->numSlots(localityDomainId)); return ret; } @@ -618,6 +799,16 @@ class GpuCacheLevelStorage : public CacheLevelStorage return mPhysMemAllocator.physMemSize(); } + int numLocalityDomains() const noexcept override + { + return mPhysMemAllocator.numLocalityDomains(); + } + + std::vector> mappedRanges(PoolGroupIndex pgIdx, PoolIndex poolIdx) const + { + return static_cast(*mPoolGroups.at(pgIdx)).mappedRanges(poolIdx); + } + private: PooledPhysMemAllocator& mPhysMemAllocator; }; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp index ccfad1c559e2..9dd1a260f703 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp @@ -22,6 +22,7 @@ #include "kv_cache_manager_v2/copyEngine.h" #include "kv_cache_manager_v2/exceptions.h" #include "kv_cache_manager_v2/kvCache.h" +#include "kv_cache_manager_v2/localityDomain.h" #include "kv_cache_manager_v2/page.h" #include "kv_cache_manager_v2/stagingBuffer.h" #include "kv_cache_manager_v2/utils/hostMem.h" @@ -266,9 +267,31 @@ StorageManager::StorageManager(LifeCycleRegistry const& lifeCycles, StorageConfi slotSizeLists.push_back(sd.slotSizeList()); } + auto const& gpuTierConfig = std::get(config.cacheTiers[kHotLevel]); size_t const gpuQuota = cacheTierQuota(config.cacheTiers[kHotLevel]); - mGpuPhysMemAllocator - = std::make_unique(tierAllocationGranularity(CacheTier::GPU_MEM, gpuQuota)); + size_t const gpuTierGranularity = tierAllocationGranularity(CacheTier::GPU_MEM, gpuQuota); + mNumLocalityDomains = resolveNumLocalityDomains(gpuTierConfig.enableLocalityDomains); + if (mNumLocalityDomains > 1) + { + // Each domain can strand up to one partial chunk per pool group. A slot size with a small power-of-two + // factor makes that chunk large; past this bound the split wastes more memory than locality is worth. + size_t chunkBytes = 0; + for (auto const& slotSizeList : slotSizeLists) + { + chunkBytes += localityDomainChunkSlots(slotSizeList, gpuTierGranularity) + * std::accumulate(slotSizeList.begin(), slotSizeList.end(), size_t{0}); + } + size_t const maxChunkBytes = std::max(gpuQuota / mNumLocalityDomains / 16, 64 * gpuTierGranularity); + if (chunkBytes > maxChunkBytes) + { + TLLM_LOG_WARNING( + "KVCacheManagerV2: locality domains disabled; the slot sizes need %zu MiB chunks per domain, above " + "the %zu MiB limit", + chunkBytes >> 20, maxChunkBytes >> 20); + mNumLocalityDomains = 1; + } + } + mGpuPhysMemAllocator = std::make_unique(gpuTierGranularity, mNumLocalityDomains); size_t const gpuGranularity = mGpuPhysMemAllocator->physMemSize(); // Constraints are hot-level feasibility floors. They are scaled by @@ -448,6 +471,15 @@ StorageManager::StorageManager(LifeCycleRegistry const& lifeCycles, StorageConfi = computeSlotCountForLevel(config.cacheTiers[level], coldSlotSizeLists, coldRatio, coldMinSlots); auto* gpuPhysMemAllocator = cacheTierOf(config.cacheTiers[level]) == CacheTier::GPU_MEM ? mGpuPhysMemAllocator.get() : nullptr; + // The shared allocator would place this tier per domain while allocation always targets + // domain 0 on non-hot levels, stranding the other domains' share. + if (gpuPhysMemAllocator != nullptr && mNumLocalityDomains > 1) + { + throw std::invalid_argument( + "Locality domains are supported only when the GPU tier is the hot tier; " + "cache tier " + + std::to_string(level.value()) + " is a GPU tier"); + } mLevels.emplace_back(lifeCycleGrouping(level), level, config.cacheTiers[level], slotDescList(level), slotCounts, gpuPhysMemAllocator, tierBackingOptions(config.cacheTiers[level])); } @@ -491,13 +523,37 @@ void StorageManager::destroy() // newSlots // --------------------------------------------------------------------------- +LocalityDomainId StorageManager::resolveLocalityDomain( + CacheLevel level, std::optional localityDomainId) const +{ + if (level != kHotLevel || mNumLocalityDomains == 1) + { + return LocalityDomainId{0}; + } + TLLM_CHECK_WITH_INFO(localityDomainId.has_value(), "locality_domain_id is required on a localized GPU tier"); + TLLM_CHECK_WITH_INFO(*localityDomainId >= 0 && *localityDomainId < mNumLocalityDomains, + "locality_domain_id %d out of range for %d locality domain(s)", localityDomainId->value(), mNumLocalityDomains); + return *localityDomainId; +} + +std::optional StorageManager::localityDomainFilter( + CacheLevel level, std::optional localityDomainId) const +{ + if (level != kHotLevel || mNumLocalityDomains == 1) + { + return std::nullopt; + } + return resolveLocalityDomain(level, localityDomainId); +} + TypedVec> StorageManager::newSlots(CacheLevel level, TypedVec const& numSlotsPerLc, MigrationRecorder const& migrationRecorder, - DropRecorder const& dropRecorder) + DropRecorder const& dropRecorder, std::optional localityDomainId) { auto const& grouping = lifeCycleGrouping(level); TLLM_CHECK_DEBUG(numSlotsPerLc.size() == numLifeCycles()); auto& storage = *mLevels.at(level).storage; + LocalityDomainId const domain = resolveLocalityDomain(level, localityDomainId); // Aggregate by pool group. TypedVec pgNumSlots(numPoolGroups(level), 0); @@ -515,7 +571,7 @@ TypedVec> StorageManager::newSlots(CacheLevel lev bool needMore = false; for (PoolGroupIndex pgIdx{0}; pgIdx < pgNumSlots.size(); ++pgIdx) { - if (pgNumSlots[pgIdx] > storage.numFreeSlots(pgIdx)) + if (pgNumSlots[pgIdx] > storage.numFreeSlots(pgIdx, domain)) { needMore = true; break; @@ -524,13 +580,13 @@ TypedVec> StorageManager::newSlots(CacheLevel lev if (needMore) { - prepareFreeSlots(level, pgNumSlots, migrationRecorder, dropRecorder); + prepareFreeSlots(level, pgNumSlots, migrationRecorder, dropRecorder, localityDomainId); } // A14: post-condition — free-slot counts satisfy requirements. for (PoolGroupIndex pgIdx{0}; pgIdx < pgNumSlots.size(); ++pgIdx) { - TLLM_CHECK_DEBUG_WITH_INFO(pgNumSlots[pgIdx] <= storage.numFreeSlots(pgIdx), + TLLM_CHECK_DEBUG_WITH_INFO(pgNumSlots[pgIdx] <= storage.numFreeSlots(pgIdx, domain), "Free slot count does not satisfy requirement after prepareFreeSlots"); } @@ -541,7 +597,7 @@ TypedVec> StorageManager::newSlots(CacheLevel lev for (LifeCycleId lcId{0}; lcId < ret.size(); ++lcId) { PoolGroupIndex pg = grouping[lcId]; - ret[lcId] = storage.allocateMultiple(pg, numSlotsPerLc[lcId]); + ret[lcId] = storage.allocateMultiple(pg, numSlotsPerLc[lcId], domain); } } catch (...) @@ -559,27 +615,29 @@ TypedVec> StorageManager::newSlots(CacheLevel lev TypedVec> StorageManager::newGpuSlots( TypedVec const& numSlotsPerLc, MigrationRecorder const& migrationRecorder, - DropRecorder const& dropRecorder) + DropRecorder const& dropRecorder, std::optional localityDomainId) { - return newSlots(kHotLevel, numSlotsPerLc, migrationRecorder, dropRecorder); + return newSlots(kHotLevel, numSlotsPerLc, migrationRecorder, dropRecorder, localityDomainId); } std::vector StorageManager::newSlotsForPoolGroup(CacheLevel level, PoolGroupIndex pgIdx, SlotCount numSlots, - MigrationRecorder const& migrationRecorder, DropRecorder const& dropRecorder) + MigrationRecorder const& migrationRecorder, DropRecorder const& dropRecorder, + std::optional localityDomainId) { if (numSlots < 0) { throw LogicError("StorageManager::newSlotsForPoolGroup: numSlots must be non-negative"); } auto& storage = *mLevels.at(level).storage; - if (numSlots > storage.numFreeSlots(pgIdx)) + LocalityDomainId const domain = resolveLocalityDomain(level, localityDomainId); + if (numSlots > storage.numFreeSlots(pgIdx, domain)) { TypedVec requirements(numPoolGroups(level), 0); requirements.at(pgIdx) = numSlots; - prepareFreeSlots(level, requirements, migrationRecorder, dropRecorder); + prepareFreeSlots(level, requirements, migrationRecorder, dropRecorder, localityDomainId); } - TLLM_CHECK_DEBUG(numSlots <= storage.numFreeSlots(pgIdx)); - return storage.allocateMultiple(pgIdx, numSlots); + TLLM_CHECK_DEBUG(numSlots <= storage.numFreeSlots(pgIdx, domain)); + return storage.allocateMultiple(pgIdx, numSlots, domain); } Address StorageManager::slotAddress(CacheLevel level, PoolGroupIndex pgIdx, SlotId slotId, PoolIndex poolIdx) const @@ -878,14 +936,16 @@ void StorageManager::excludeFromEviction(Page& page) // --------------------------------------------------------------------------- void StorageManager::prepareFreeSlots(CacheLevel level, TypedVec const& requirements, - MigrationRecorder const& migrationRecorder, DropRecorder const& dropRecorder) + MigrationRecorder const& migrationRecorder, DropRecorder const& dropRecorder, + std::optional localityDomainId) { refreshRetention(); - _prepareFreeSlots(level, requirements, migrationRecorder, dropRecorder); + _prepareFreeSlots(level, requirements, migrationRecorder, dropRecorder, localityDomainId); } void StorageManager::_prepareFreeSlots(CacheLevel level, TypedVec const& requirements, - MigrationRecorder const& migrationRecorder, DropRecorder const& dropRecorder) + MigrationRecorder const& migrationRecorder, DropRecorder const& dropRecorder, + std::optional localityDomainId) { TypedVec> goals(numCacheLevels()); for (CacheLevel lvl{0}; lvl < goals.size(); ++lvl) @@ -898,13 +958,13 @@ void StorageManager::_prepareFreeSlots(CacheLevel level, TypedVec const& minNumPages, DropRecorder const& dropRecorder) +void StorageManager::forceEvict(CacheLevel level, TypedVec const& minNumPages, + DropRecorder const& dropRecorder, std::optional localityDomainId) { - auto evicted = mLevels.at(level).controller.evict(minNumPages); + auto evicted = mLevels.at(level).controller.evict(minNumPages, localityDomainFilter(level, localityDomainId)); auto rescheduleEvictedPagesOnFailure = makeEvictionRollbackGuard(evicted); if (isLastLevel(level)) @@ -958,7 +1018,7 @@ void StorageManager::forceEvict( void StorageManager::_prepareFreeSlots(TypedVec>& goals, CacheLevel lvlId, PagesByLifeCycle& fallenPages, MigrationRecorder const& migrationRecorder, - DropRecorder const& dropRecorder) + DropRecorder const& dropRecorder, std::optional localityDomainId) { TypedVec>> belowOffloadThreshold(numLifeCycles()); auto restoreBelowThreshold = FuncGuard( @@ -1020,6 +1080,8 @@ void StorageManager::_prepareFreeSlots(TypedVec const domainFilter = localityDomainFilter(lvlId, localityDomainId); auto countPages = [&](PoolGroupIndex pgIdx, PagesByLifeCycle const& pagesByLifeCycle) { @@ -1036,8 +1098,8 @@ void StorageManager::_prepareFreeSlots(TypedVec oldFree ? required - oldFree : 0; numToEvict.at(pgIdx) = std::min(shortage, evictableCount); @@ -1074,7 +1136,7 @@ void StorageManager::_prepareFreeSlots(TypedVec>> acceptedPages(numLifeCycles()); @@ -1116,7 +1178,7 @@ void StorageManager::_prepareFreeSlots(TypedVec= numEvicted + oldFree); TLLM_CHECK_DEBUG_WITH_INFO( countPages(pgIdx, heldPages) <= freeAfterEviction, "held_pages count exceeds new free slot count"); @@ -1175,7 +1237,7 @@ void StorageManager::_prepareFreeSlots(TypedVecstatus() == PageStatus::DROPPABLE) @@ -1244,7 +1307,7 @@ LayerGroupId StorageManager::getMigrationBatchingLayerGroupId( std::optional> StorageManager::_batchedMigrate(CacheLevel dstLevel, CacheLevel srcLevel, std::vector> const& srcPages, bool updateSrc, MigrationRecorder const& migrationRecorder, - bool defrag) + bool defrag, std::optional dstLocalityDomainId) { TLLM_CHECK_DEBUG(defrag || dstLevel != srcLevel); if (srcPages.empty()) @@ -1275,11 +1338,12 @@ std::optional> StorageManager::_batchedMigrate(CacheLevel dstL })); auto& srcPoolGroup = poolGroup(srcLevel, srcPgIdx); auto& dstPoolGroup = poolGroup(dstLevel, dstPgIdx); + LocalityDomainId const dstDomain = resolveLocalityDomain(dstLevel, dstLocalityDomainId); - if (dstPoolGroup.numFreeSlots() < numSlots) + if (dstPoolGroup.numFreeSlots(dstDomain) < numSlots) throw OutOfPagesError("Not enough free slots for migration"); - auto dstSlots = dstPoolGroup.allocateMultiple(numSlots); + auto dstSlots = dstPoolGroup.allocateMultiple(numSlots, dstDomain); // A15: allocated slot count must match the request. TLLM_CHECK_DEBUG_WITH_INFO(slotCountValueFromSize(dstSlots.size()) == numSlots, "dst_slots size mismatch"); try @@ -1375,8 +1439,8 @@ std::optional> StorageManager::_batchedMigrate(CacheLevel dstL // batchedMigrate // --------------------------------------------------------------------------- -void StorageManager::batchedMigrate( - CacheLevel dstLevel, std::vector> const& pages, MigrationRecorder const& migrationRecorder) +void StorageManager::batchedMigrate(CacheLevel dstLevel, std::vector> const& pages, + MigrationRecorder const& migrationRecorder, std::optional dstLocalityDomainId) { std::map>> groups; std::set seen; @@ -1391,7 +1455,8 @@ void StorageManager::batchedMigrate( } for (auto& [key, group] : groups) { - _batchedMigrate(dstLevel, key.first, group, /*updateSrc=*/true, migrationRecorder); + _batchedMigrate( + dstLevel, key.first, group, /*updateSrc=*/true, migrationRecorder, /*defrag=*/false, dstLocalityDomainId); } } @@ -1402,19 +1467,21 @@ void StorageManager::offloadSparsePages(KvCache& requestingCache, std::vector> const& pages, - MigrationRecorder const& migrationRecorder, DropRecorder const& dropRecorder) + MigrationRecorder const& migrationRecorder, DropRecorder const& dropRecorder, + std::optional localityDomainId) { - _migrateLockedSparsePages(stream, kHotLevel, pages, nullptr, migrationRecorder, dropRecorder); + _migrateLockedSparsePages(stream, kHotLevel, pages, nullptr, migrationRecorder, dropRecorder, localityDomainId); } void StorageManager::_migrateLockedSparsePages(CUstream stream, CacheLevel dstLevel, std::vector> const& pages, KvCache const* offloadingCache, - MigrationRecorder const& migrationRecorder, DropRecorder const& dropRecorder) + MigrationRecorder const& migrationRecorder, DropRecorder const& dropRecorder, + std::optional localityDomainId) { CacheLevel const srcLevel = dstLevel == kHotLevel ? kSparseHistoryLevel : kHotLevel; @@ -1475,7 +1542,8 @@ void StorageManager::_migrateLockedSparsePages(CUstream stream, CacheLevel dstLe { requirements[getPoolGroupIndex(dstLevel, layerGroup)] += slotCountValueFromSize(batch.srcPages.size()); } - prepareFreeSlots(dstLevel, requirements, migrationRecorder, dropRecorder); + prepareFreeSlots(dstLevel, requirements, migrationRecorder, dropRecorder, localityDomainId); + LocalityDomainId const dstDomain = resolveLocalityDomain(dstLevel, localityDomainId); auto releaseDestinations = FuncGuard( [&]() { @@ -1493,7 +1561,7 @@ void StorageManager::_migrateLockedSparsePages(CUstream stream, CacheLevel dstLe for (auto& [layerGroup, batch] : batches) { auto& pool = poolGroup(dstLevel, getPoolGroupIndex(dstLevel, layerGroup)); - batch.dstSlots = pool.allocateMultiple(slotCountValueFromSize(batch.srcPages.size())); + batch.dstSlots = pool.allocateMultiple(slotCountValueFromSize(batch.srcPages.size()), dstDomain); batch.srcDstPageIndices.reserve(batch.srcPages.size()); for (size_t i = 0; i < batch.srcPages.size(); ++i) { @@ -1585,8 +1653,9 @@ void StorageManager::_migrateLockedSparsePages(CUstream stream, CacheLevel dstLe } } -int64_t StorageManager::prefetch( - CacheLevel dstLevel, TypedVec>>> const& pages) +int64_t StorageManager::prefetch(CacheLevel dstLevel, + TypedVec>>> const& pages, + std::optional dstLocalityDomainId) { TypedVec numSlotsToMigrate(numPoolGroups(dstLevel), 0); int64_t diskBlocksMigrated = 0; @@ -1631,7 +1700,7 @@ int64_t StorageManager::prefetch( } } - prepareFreeSlots(dstLevel, numSlotsToMigrate); + prepareFreeSlots(dstLevel, numSlotsToMigrate, {}, {}, dstLocalityDomainId); std::map>> migrationGroups; for (LifeCycleId lifeCycle{0}; lifeCycle < pages.size(); ++lifeCycle) { @@ -1664,11 +1733,12 @@ int64_t StorageManager::prefetch( { TemporaryCudaStream stream({}); auto scope = stream.enter(); - promoteSparsePages(stream.get(), lockedPages); + promoteSparsePages(stream.get(), lockedPages, {}, {}, dstLocalityDomainId); std::erase_if(migrationPages, [](auto const& page) { return page->cacheLevel == kHotLevel; }); } } - _batchedMigrate(dstLevel, srcLevel, migrationPages, /*updateSrc=*/true); + _batchedMigrate( + dstLevel, srcLevel, migrationPages, /*updateSrc=*/true, {}, /*defrag=*/false, dstLocalityDomainId); // Per batch, after it landed: the pages are already grouped by source level, so this costs // nothing per page and never credits a batch that did not run. if (cacheTier(srcLevel) == CacheTier::DISK) @@ -1749,6 +1819,7 @@ MemAddress StorageManager::getMemPoolBaseAddress(PoolGroupIndex pgIdx, PoolIndex TypedVec StorageManager::poolGroupDescs() const { auto const& descList = slotDescList(kHotLevel); + auto const& gpuStorage = static_cast(*mLevels[kHotLevel].storage); TypedVec result; result.reserve(descList.size()); for (PoolGroupIndex pgIdx{0}; pgIdx < descList.size(); ++pgIdx) @@ -1758,9 +1829,12 @@ TypedVec StorageManager::poolGroupDescs() const pools.reserve(slotSizeList.size()); for (PoolIndex poolIdx{0}; poolIdx < slotSizeList.size(); ++poolIdx) { - pools.push_back(PoolDesc{poolIdx, getMemPoolBaseAddress(pgIdx, poolIdx), slotSizeList.at(poolIdx)}); + pools.push_back(PoolDesc{poolIdx, getMemPoolBaseAddress(pgIdx, poolIdx), slotSizeList.at(poolIdx), + gpuStorage.mappedRanges(pgIdx, poolIdx)}); } - result.push_back(PoolGroupDesc{pgIdx, numSlots(pgIdx, kHotLevel), descList.at(pgIdx), std::move(pools)}); + // Localized groups have unmapped holes between domains, so the slot count is the index bound. + result.push_back( + PoolGroupDesc{pgIdx, gpuStorage.slotIndexUpperBound(pgIdx), descList.at(pgIdx), std::move(pools)}); } return result; } @@ -1778,23 +1852,44 @@ SlotCount StorageManager::numSlots(PoolGroupIndex pgIdx, CacheLevel level) const return mLevels.at(level).storage->numSlots(pgIdx); } -StorageStatistics StorageManager::getStatistics(CacheLevel level, PoolGroupIndex pgIdx) const +SlotCount StorageManager::slotIndexUpperBound(PoolGroupIndex pgIdx, CacheLevel level) const +{ + return mLevels.at(level).storage->slotIndexUpperBound(pgIdx); +} + +std::vector> StorageManager::slotRanges(PoolGroupIndex pgIdx, CacheLevel level) const +{ + return mLevels.at(level).storage->slotRanges(pgIdx); +} + +StorageStatistics StorageManager::getStatistics( + CacheLevel level, PoolGroupIndex pgIdx, std::optional localityDomainId) const { auto const& lvl = mLevels.at(level); - SlotCount freeSlots = lvl.storage->numFreeSlots(pgIdx); - SlotCount totalSlots = lvl.storage->numSlots(pgIdx); - SlotCount evictable = lvl.controller.numEvictablePages(pgIdx); + auto const domainFilter + = localityDomainId.has_value() ? localityDomainFilter(level, localityDomainId) : std::nullopt; + SlotCount freeSlots = 0; + SlotCount totalSlots = 0; + for (LocalityDomainId domain{0}; domain < lvl.storage->numLocalityDomains(); ++domain) + { + if (domainFilter.has_value() && domain != *domainFilter) + continue; + freeSlots += lvl.storage->numFreeSlots(pgIdx, domain); + totalSlots += lvl.storage->numSlots(pgIdx, domain); + } + SlotCount evictable = lvl.controller.numEvictablePages(pgIdx, domainFilter); auto sizes = lvl.storage->slotSize(pgIdx); return StorageStatistics{sizes, totalSlots, freeSlots, evictable}; } -TypedVec StorageManager::getUtilization(CacheLevel level) const +TypedVec StorageManager::getUtilization( + CacheLevel level, std::optional localityDomainId) const { TypedVec result; result.reserve(numPoolGroups(level)); for (PoolGroupIndex pgIdx{0}; pgIdx < numPoolGroups(level); ++pgIdx) { - auto const s = getStatistics(level, pgIdx); + auto const s = getStatistics(level, pgIdx, localityDomainId); TLLM_CHECK_WITH_INFO(s.total > 0, "getUtilization: pool group %d at level %d has zero capacity", static_cast(pgIdx.value()), static_cast(level.value())); result.push_back(static_cast(s.unavailable()) / static_cast(s.total)); @@ -1822,26 +1917,42 @@ float StorageManager::getOverallUtilization(CacheLevel level) const // expandPoolGroup // --------------------------------------------------------------------------- -void StorageManager::expandPoolGroup(CacheLevel level, PoolGroupIndex pgIdx, SlotCount newNumSlots) +void StorageManager::expandPoolGroup( + CacheLevel level, PoolGroupIndex pgIdx, SlotCount newNumSlots, LocalityDomainId localityDomainId) { auto& pg = poolGroup(level, pgIdx); - TLLM_CHECK_DEBUG(newNumSlots > pg.numSlots()); - pg.resizePools(newNumSlots); - pg.slotAllocator().expand(newNumSlots); + TLLM_CHECK_WITH_INFO(newNumSlots > pg.numSlots(localityDomainId), + "expandPoolGroup: target %d is not above the current %d slots of locality domain %d", + static_cast(newNumSlots), static_cast(pg.numSlots(localityDomainId)), localityDomainId.value()); + pg.resizePools(newNumSlots, localityDomainId); + pg.slotAllocator(localityDomainId).expand(newNumSlots); } // --------------------------------------------------------------------------- // shrinkPoolGroup — mirrors Python _storage_manager.py::shrink_pool_group // --------------------------------------------------------------------------- -void StorageManager::shrinkPoolGroup( - CacheLevel level, PoolGroupIndex pgIdx, SlotCount newNumSlots, std::vector> const& persistentPages) +void StorageManager::shrinkPoolGroup(CacheLevel level, PoolGroupIndex pgIdx, SlotCount newNumSlots, + std::vector> const& persistentPages, LocalityDomainId localityDomainId) { refreshRetention(); auto& pg = poolGroup(level, pgIdx); - auto& allocator = pg.slotAllocator(); + auto& allocator = pg.slotAllocator(localityDomainId); auto& ctrl = mLevels.at(level).controller; - TLLM_CHECK_DEBUG(newNumSlots < pg.numSlots()); + TLLM_CHECK_WITH_INFO(newNumSlots < pg.numSlots(localityDomainId), + "shrinkPoolGroup: target %d is not below the current %d slots of locality domain %d", + static_cast(newNumSlots), static_cast(pg.numSlots(localityDomainId)), localityDomainId.value()); + // Only this domain's pages and local slot indices matter; slot ids are global and interleaved. + std::optional const domainFilter + = pg.numLocalityDomains() == 1 ? std::nullopt : std::optional{localityDomainId}; + auto const inScope + = [&](Page const& page) { return !domainFilter.has_value() || page.localityDomainId() == domainFilter; }; + auto const overflows = [&](Page const& page) + { + auto const local = allocator.localIndexOf(page.slotId()); + TLLM_CHECK_DEBUG(local.has_value()); + return local.value_or(SlotCount{0}) >= newNumSlots; + }; // A16: persistent_pages preconditions. TLLM_CHECK_DEBUG_WITH_INFO( @@ -1858,18 +1969,20 @@ void StorageManager::shrinkPoolGroup( { allocator.prepareForShrink(newNumSlots); allocator.finishShrink(); - pg.resizePools(newNumSlots); + pg.resizePools(newNumSlots, localityDomainId); return; } - // Find overflow pages: scheduled pages with slot_id >= newNumSlots. + // Find overflow pages: scheduled pages of this domain whose local index >= newNumSlots. auto gen = ctrl.pageGenerator(pgIdx); std::deque>> overflowSlots; { SlotCount idx = 0; while (auto const* page = gen()) { - if ((*page)->slotId() >= newNumSlots) + if (!inScope(**page)) + continue; + if (overflows(**page)) overflowSlots.emplace_back(idx, *page); ++idx; } @@ -1879,7 +1992,7 @@ void StorageManager::shrinkPoolGroup( std::vector> overflowPersistent; for (auto const& p : persistentPages) { - if (p->slotId() >= newNumSlots) + if (inScope(*p) && overflows(*p)) overflowPersistent.push_back(p); } SlotCount numOverflowPersistent = slotCountValueFromSize(overflowPersistent.size()); @@ -1910,7 +2023,7 @@ void StorageManager::shrinkPoolGroup( // Force-evict the required pages. TypedVec evictReqs(numPoolGroups(level), 0); evictReqs[pgIdx] = minNumEvicted; - forceEvict(level, evictReqs); + forceEvict(level, evictReqs, {}, domainFilter); // Remaining overflow pages to defragment. std::vector> overflowPages; @@ -1924,15 +2037,16 @@ void StorageManager::shrinkPoolGroup( TypedVec reqs(numPoolGroups(level), 0); reqs[pgIdx] = slotCountValueFromSize(overflowPages.size()); // The eviction order captured above must remain stable through defragmentation. - _prepareFreeSlots(level, reqs, MigrationRecorder{}, DropRecorder{}); + _prepareFreeSlots(level, reqs, MigrationRecorder{}, DropRecorder{}, domainFilter); // A17: all overflow pages must be at the expected cache level. TLLM_CHECK_DEBUG_WITH_INFO(std::all_of(overflowPages.begin(), overflowPages.end(), [level](auto const& p) { return p->cacheLevel == level; }), "Overflow page cache level mismatch"); - // Defragment: migrate overflow pages to free slots within the same level. - _batchedMigrate(level, level, overflowPages, /*updateSrc=*/true, MigrationRecorder{}, /*defrag=*/true); + // Defragment: migrate overflow pages to free slots within the same level and domain. + _batchedMigrate( + level, level, overflowPages, /*updateSrc=*/true, MigrationRecorder{}, /*defrag=*/true, domainFilter); // A18: post-defrag overflow assertion — overflow slot count matches expectations. TLLM_CHECK_DEBUG_WITH_INFO(allocator.numOverflowSlots() == allocator.numActiveSlots() - allocator.targetCapacity(), @@ -1940,7 +2054,7 @@ void StorageManager::shrinkPoolGroup( // Finalize shrink and resize pools. allocator.finishShrink(); - pg.resizePools(newNumSlots); + pg.resizePools(newNumSlots, localityDomainId); } // --------------------------------------------------------------------------- @@ -1952,10 +2066,19 @@ void StorageManager::adjustCacheLevel(CacheLevel level, std::optional ne TypedVec>> const* persistentPages) { auto& lvlStorage = *mLevels.at(level).storage; - auto oldNumSlots = lvlStorage.slotCountList(); - size_t quota = newQuota.has_value() - ? roundUp(newQuota.value(), static_cast(lvlStorage.poolSizeGranularity())) - : lvlStorage.totalQuota(); + size_t const granularity = lvlStorage.poolSizeGranularity(); + int const numDomains = lvlStorage.numLocalityDomains(); + size_t quota; + if (numDomains > 1) + { + // Slot counts and the feasibility floor are per domain. The share is taken from the + // requested quota before any rounding, as construction does, so init and resize agree. + quota = roundDown(newQuota.value_or(lvlStorage.totalQuota()) / static_cast(numDomains), granularity); + } + else + { + quota = newQuota.has_value() ? roundUp(newQuota.value(), granularity) : lvlStorage.totalQuota(); + } auto const minSlots = level == kHotLevel ? mMinSlots : TypedVec(numPoolGroups(level), SlotCount{1}); size_t minQuota = minQuotaForLevel(lvlStorage.slotSizeLists(), lvlStorage.poolSizeGranularity(), minSlots); @@ -1976,22 +2099,42 @@ void StorageManager::adjustCacheLevel(CacheLevel level, std::optional ne TLLM_CHECK_DEBUG(isLastLevel(level) || persistentPages == nullptr); + // Domains are resized one at a time and a failure leaves the already resized ones at the new + // capacity, so each (pool group, domain) pair is compared against its own current capacity: + // a retry then only touches the pairs that are still off target. // Shrink first. for (PoolGroupIndex pgIdx{0}; pgIdx < newNumSlots.size(); ++pgIdx) { - if (newNumSlots[pgIdx] >= oldNumSlots[pgIdx]) - continue; std::vector> pages; if (persistentPages) pages = (*persistentPages)[pgIdx]; - shrinkPoolGroup(level, pgIdx, newNumSlots[pgIdx], pages); + for (LocalityDomainId localityDomainId{0}; localityDomainId < numDomains; ++localityDomainId) + { + if (newNumSlots[pgIdx] >= lvlStorage.numSlots(pgIdx, localityDomainId)) + continue; + if (numDomains == 1) + { + shrinkPoolGroup(level, pgIdx, newNumSlots[pgIdx], pages, localityDomainId); + continue; + } + // The persistent list spans every domain, but each shrink only keeps its own. + std::vector> domainPages; + for (auto const& page : pages) + { + if (page->localityDomainId() == std::optional{localityDomainId}) + domainPages.push_back(page); + } + shrinkPoolGroup(level, pgIdx, newNumSlots[pgIdx], domainPages, localityDomainId); + } } // Then expand. for (PoolGroupIndex pgIdx{0}; pgIdx < newNumSlots.size(); ++pgIdx) { - if (newNumSlots[pgIdx] <= oldNumSlots[pgIdx]) - continue; - expandPoolGroup(level, pgIdx, newNumSlots[pgIdx]); + for (LocalityDomainId localityDomainId{0}; localityDomainId < numDomains; ++localityDomainId) + { + if (newNumSlots[pgIdx] > lvlStorage.numSlots(pgIdx, localityDomainId)) + expandPoolGroup(level, pgIdx, newNumSlots[pgIdx], localityDomainId); + } } if (lvlStorage.cacheTier() == CacheTier::GPU_MEM) { @@ -2313,6 +2456,12 @@ TypedVec StorageManager::computeSlotCountForLevel(Cac { size_t quota = cacheTierQuota(tierConfig); size_t const granularity = tierCommitUnit(tierConfig); + if (cacheTierOf(tierConfig) == CacheTier::GPU_MEM && mNumLocalityDomains > 1) + { + // Each domain rounds its own backing up to the granularity, so the slot count is sized + // per domain from an equal share of the quota rather than split after sizing. + quota = roundDown(quota / static_cast(mNumLocalityDomains), granularity); + } quota = std::max(minQuotaForLevel(slotSizeLists, granularity, minSlots), roundUp(quota, granularity)); return CacheLevelStorage::ratioToSlotCountList(quota, slotSizeLists, ratio, granularity, minSlots); } diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.h index 2b423590aeb7..4a3b2d240523 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.h @@ -143,17 +143,21 @@ class StorageManager : public std::enable_shared_from_this // Allocate slots for all life cycles at the given cache level. // numSlotsPerLc[lcId] = how many slots to allocate for that life cycle. // Returns a vector indexed by lcId. + // localityDomainId selects the GPU domain the slots come from; it is required when the GPU + // tier is localized and ignored for the other tiers. TypedVec> newSlots(CacheLevel level, TypedVec const& numSlotsPerLc, MigrationRecorder const& migrationRecorder = {}, - DropRecorder const& dropRecorder = {}); + DropRecorder const& dropRecorder = {}, std::optional localityDomainId = std::nullopt); TypedVec> newGpuSlots(TypedVec const& numSlotsPerLc, - MigrationRecorder const& migrationRecorder = {}, DropRecorder const& dropRecorder = {}); + MigrationRecorder const& migrationRecorder = {}, DropRecorder const& dropRecorder = {}, + std::optional localityDomainId = std::nullopt); // Allocate slots for a single pool group at the given cache level. // Returns numSlots Slot objects. Throws OutOfPagesError if allocation fails. std::vector newSlotsForPoolGroup(CacheLevel level, PoolGroupIndex pgIdx, SlotCount numSlots, - MigrationRecorder const& migrationRecorder = {}, DropRecorder const& dropRecorder = {}); + MigrationRecorder const& migrationRecorder = {}, DropRecorder const& dropRecorder = {}, + std::optional localityDomainId = std::nullopt); // Release a slot back to its pool. void releaseSlot(LifeCycleId lc, CacheLevel level, Slot slot); @@ -175,27 +179,31 @@ class StorageManager : public std::enable_shared_from_this // Check if a page is evictable (optionally at a target level). bool isEvictable(Page const& page, std::optional level = std::nullopt) const noexcept; - // Ensure the requested free slots exist at this level (evicting pages as needed). + // Ensure the requested free slots exist at this level (evicting pages as needed). On a + // localized GPU tier only pages of localityDomainId are evicted and counted as free. void prepareFreeSlots(CacheLevel level, TypedVec const& requirements, - MigrationRecorder const& migrationRecorder = {}, DropRecorder const& dropRecorder = {}); + MigrationRecorder const& migrationRecorder = {}, DropRecorder const& dropRecorder = {}, + std::optional localityDomainId = std::nullopt); // Force-evict pages from a level to free space. void forceEvict(CacheLevel level, TypedVec const& minNumPages, - DropRecorder const& dropRecorder = {}); + DropRecorder const& dropRecorder = {}, std::optional localityDomainId = std::nullopt); - // Dynamic cache level resizing. + // Dynamic cache level resizing. A localized GPU level is resized one domain at a time. void adjustCacheLevel(CacheLevel level, std::optional newQuota, TypedVec const& ratioList, TypedVec>> const* persistentPages); void shrinkPoolGroup(CacheLevel level, PoolGroupIndex pgIdx, SlotCount newNumSlots, - std::vector> const& persistentPages); - void expandPoolGroup(CacheLevel level, PoolGroupIndex pgIdx, SlotCount newNumSlots); + std::vector> const& persistentPages, LocalityDomainId localityDomainId = LocalityDomainId{0}); + void expandPoolGroup(CacheLevel level, PoolGroupIndex pgIdx, SlotCount newNumSlots, + LocalityDomainId localityDomainId = LocalityDomainId{0}); // ---- Migration --------------------------------------------------------- // Migrate pages excluded from eviction to one destination level. Locked pages cannot move. - void batchedMigrate( - CacheLevel dstLevel, std::vector> const& pages, MigrationRecorder const& migrationRecorder); + // dstLocalityDomainId is required when the destination is a localized GPU tier. + void batchedMigrate(CacheLevel dstLevel, std::vector> const& pages, + MigrationRecorder const& migrationRecorder, std::optional dstLocalityDomainId = std::nullopt); //! Demote complete locked sparse pages on the requesting owner's stream, updating every live owner. //! Caller holds the manager's exclusive lock. Allocation/copy failure preserves GPU ownership. @@ -204,13 +212,17 @@ class StorageManager : public std::enable_shared_from_this //! Promote locked sparse host pages into one shared GPU allocation per page, updating all owners. //! Caller holds the manager's exclusive lock. Allocation/copy failure preserves host ownership. + // localityDomainId is required when the GPU tier is localized: the promoted pages land there. void promoteSparsePages(CUstream stream, std::vector> const& pages, - MigrationRecorder const& migrationRecorder = {}, DropRecorder const& dropRecorder = {}); + MigrationRecorder const& migrationRecorder = {}, DropRecorder const& dropRecorder = {}, + std::optional localityDomainId = std::nullopt); // Best-effort migration of grouped pages to a destination cache level. Returns how many pages // it moved off the disk tier. If a later batch fails, earlier batches may already have migrated. - int64_t prefetch( - CacheLevel dstLevel, TypedVec>>> const& pages); + // dstLocalityDomainId is required when the destination is a localized GPU tier. + int64_t prefetch(CacheLevel dstLevel, + TypedVec>>> const& pages, + std::optional dstLocalityDomainId = std::nullopt); // ---- Query helpers ----------------------------------------------------- @@ -231,6 +243,12 @@ class StorageManager : public std::enable_shared_from_this return mSlotDescLists.at(level).size(); } + // 1 unless the GPU tier is localized. + int numLocalityDomains() const noexcept + { + return mNumLocalityDomains; + } + PoolGroupIndex numPoolGroups() const { return numPoolGroups(kHotLevel); @@ -318,12 +336,19 @@ class StorageManager : public std::enable_shared_from_this // NOTE: Python's get_statistics(level) returns a list over all pool groups. // C++ takes a single pgIdx for flexibility; the nanobind wrapper loops over // all pool groups to match Python's signature. - StorageStatistics getStatistics(CacheLevel level = kHotLevel, PoolGroupIndex pgIdx = PoolGroupIndex{0}) const; - TypedVec getUtilization(CacheLevel level = kHotLevel) const; + // Summed over locality domains unless one is given. + StorageStatistics getStatistics(CacheLevel level = kHotLevel, PoolGroupIndex pgIdx = PoolGroupIndex{0}, + std::optional localityDomainId = std::nullopt) const; + TypedVec getUtilization( + CacheLevel level = kHotLevel, std::optional localityDomainId = std::nullopt) const; float getOverallUtilization(CacheLevel level = kHotLevel) const; // Pool-group slot count (number of pages). + // Slots of locality domain 0 (the whole group when not localized). SlotCount numSlots(PoolGroupIndex pgIdx, CacheLevel level = kHotLevel) const; + // One past the highest usable slot id across all locality domains. + SlotCount slotIndexUpperBound(PoolGroupIndex pgIdx, CacheLevel level = kHotLevel) const; + std::vector> slotRanges(PoolGroupIndex pgIdx, CacheLevel level = kHotLevel) const; // Layer-to-lifecycle mapping (for KvCacheManager queries). std::unordered_map const& layerToLifeCycleIds() const noexcept @@ -396,16 +421,24 @@ class StorageManager : public std::enable_shared_from_this [[nodiscard]] auto makeEvictionRollbackGuard(TypedVec>> const& evicted); void _prepareFreeSlots(CacheLevel level, TypedVec const& requirements, - MigrationRecorder const& migrationRecorder, DropRecorder const& dropRecorder); + MigrationRecorder const& migrationRecorder, DropRecorder const& dropRecorder, + std::optional localityDomainId = std::nullopt); void _prepareFreeSlots(TypedVec>& goals, CacheLevel lvlId, PagesByLifeCycle& fallenPages, MigrationRecorder const& migrationRecorder = {}, - DropRecorder const& dropRecorder = {}); + DropRecorder const& dropRecorder = {}, std::optional localityDomainId = std::nullopt); std::optional> _batchedMigrate(CacheLevel dstLevel, CacheLevel srcLevel, std::vector> const& srcPages, bool updateSrc, MigrationRecorder const& migrationRecorder = {}, - bool defrag = false); + bool defrag = false, std::optional dstLocalityDomainId = std::nullopt); void _migrateLockedSparsePages(CUstream stream, CacheLevel dstLevel, std::vector> const& pages, - KvCache const* offloadingCache, MigrationRecorder const& migrationRecorder, DropRecorder const& dropRecorder); + KvCache const* offloadingCache, MigrationRecorder const& migrationRecorder, DropRecorder const& dropRecorder, + std::optional localityDomainId); + + // The domain whose allocator serves `level`: 0 unless the level is a localized GPU tier, which + // then requires an explicit domain. Also the eviction filter for that level, if any. + LocalityDomainId resolveLocalityDomain(CacheLevel level, std::optional localityDomainId) const; + std::optional localityDomainFilter( + CacheLevel level, std::optional localityDomainId) const; [[nodiscard]] LayerGroupId getMigrationBatchingLayerGroupId( CacheLevel dstLevel, CacheLevel srcLevel, LifeCycleId lifeCycle) const; void submitMigrationBatch(CacheLevel dstLevel, CacheLevel srcLevel, LayerGroupId batchingLayerGroupId, @@ -474,6 +507,7 @@ class StorageManager : public std::enable_shared_from_this // All GPU cache levels borrow this allocator. It must outlive mLevels. std::unique_ptr mGpuPhysMemAllocator; RetentionExpiryQueue mRetentionExpiries; + int mNumLocalityDomains{1}; TypedVec mLevels; static constexpr size_t kDefaultPageStagingBytes = 64u << 20u; diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index d56f2bb5fc83..7ce9292a4e12 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -367,8 +367,13 @@ static nb::object optionalIntToObject(std::optional value) return nb::cast(*value); } +// Two fields unless a locality domain is set, so non-localized scopes keep their tuple shape. static nb::tuple reuseScopeTuple(kv::ReuseScope const& self) { + if (self.localityDomainId.has_value()) + { + return nb::make_tuple(optionalIntToObject(self.loraId), self.salt, nb::cast(self.localityDomainId->value())); + } return nb::make_tuple(optionalIntToObject(self.loraId), self.salt); } @@ -710,8 +715,17 @@ static kv::ReuseScope castReuseScope(nb::object reuseScope) } if (PyObject_HasAttrString(reuseScope.ptr(), "lora_id") && PyObject_HasAttrString(reuseScope.ptr(), "salt")) { + std::optional localityDomainId; + if (PyObject_HasAttrString(reuseScope.ptr(), "locality_domain_id")) + { + auto const value = castOptionalIntAttr(reuseScope, "locality_domain_id"); + if (value.has_value()) + { + localityDomainId = kv::LocalityDomainId{static_cast(*value)}; + } + } return {castOptionalIntAttr(reuseScope, "lora_id"), - nb::cast>(nb::getattr(reuseScope, "salt"))}; + nb::cast>(nb::getattr(reuseScope, "salt")), localityDomainId}; } throw std::invalid_argument( "reuse_scope must be None, ReuseScope, an int lora_task_id, or an object with lora_id and salt"); @@ -1433,30 +1447,34 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) nb::class_(m, "ReuseScope") .def( "__init__", - [](kv::ReuseScope* self, std::optional loraId, std::optional salt) { - new (self) kv::ReuseScope{loraId, std::move(salt)}; + [](kv::ReuseScope* self, std::optional loraId, std::optional salt, + std::optional localityDomainId) { + new (self) + kv::ReuseScope{loraId, std::move(salt), std::optional(localityDomainId)}; }, - nb::arg("lora_id").none() = std::nullopt, nb::arg("salt").none() = std::nullopt) + nb::arg("lora_id").none() = std::nullopt, nb::arg("salt").none() = std::nullopt, + nb::arg("locality_domain_id").none() = std::nullopt) .def_ro("lora_id", &kv::ReuseScope::loraId) .def_ro("salt", &kv::ReuseScope::salt) - .def("__len__", [](kv::ReuseScope const&) { return 2; }) + .def_prop_ro("locality_domain_id", + [](kv::ReuseScope const& self) + { return self.localityDomainId ? std::optional{self.localityDomainId->value()} : std::nullopt; }) + .def("__len__", [](kv::ReuseScope const& self) { return reuseScopeTuple(self).size(); }) .def( "__getitem__", [](kv::ReuseScope const& self, int index) -> nb::object { + auto const fields = reuseScopeTuple(self); + auto const size = static_cast(fields.size()); if (index < 0) { - index += 2; + index += size; } - if (index == 0) - { - return optionalIntToObject(self.loraId); - } - if (index == 1) + if (index < 0 || index >= size) { - return nb::cast(self.salt); + throw nb::index_error("ReuseScope index out of range"); } - throw nb::index_error("ReuseScope index out of range"); + return fields[index]; }, nb::arg("index")) .def("__iter__", @@ -1471,7 +1489,7 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) { return nb::bool_(self == nb::cast(other)); } - if (PyTuple_Check(other.ptr()) && PyTuple_GET_SIZE(other.ptr()) == 2) + if (PyTuple_Check(other.ptr())) { return nb::bool_(PyObject_RichCompareBool(reuseScopeTuple(self).ptr(), other.ptr(), Py_EQ) == 1); } @@ -1485,6 +1503,10 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) repr += self.loraId.has_value() ? std::to_string(*self.loraId) : "None"; repr += ", salt="; repr += pythonRepr(self.salt); + if (self.localityDomainId.has_value()) + { + repr += ", locality_domain_id=" + std::to_string(self.localityDomainId->value()); + } repr += ")"; return repr; }); @@ -1563,6 +1585,7 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) nb::class_(m, "PoolDesc") .def_prop_ro("pool_index", [](kv::PoolDesc const& self) { return self.poolIndex.value(); }) .def_ro("base_address", &kv::PoolDesc::baseAddress) + .def_ro("mapped_ranges", &kv::PoolDesc::mappedRanges) .def_ro("slot_bytes", &kv::PoolDesc::slotBytes); nb::class_(m, "PoolGroupDesc") @@ -1594,8 +1617,14 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) "__deepcopy__", [](cls const& self, nb::dict) { return cls(self); }, nb::arg("memo")) nb::class_(m, "GpuCacheTierConfig") - .def(nb::init(), nb::arg("quota")) + .def( + "__init__", + [](kv::GpuCacheTierConfig* self, size_t quota, bool enableLocalityDomains) { + new (self) kv::GpuCacheTierConfig{quota, enableLocalityDomains}; + }, + nb::arg("quota"), nb::arg("enable_locality_domains") = false) .def_rw("quota", &kv::GpuCacheTierConfig::quota) + .def_rw("enable_locality_domains", &kv::GpuCacheTierConfig::enableLocalityDomains) .def_prop_ro("tier", &kv::GpuCacheTierConfig::tier) .def("validate", &kv::GpuCacheTierConfig::validate) DEF_COPY(kv::GpuCacheTierConfig); @@ -2112,6 +2141,12 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) nb::rv_policy::reference_internal) .def_prop_ro("committed_tokens", &committedTokensList) .def_prop_ro("reuse_scope", &kv::KvCache::reuseScope) + .def_prop_ro("locality_domain_id", + [](kv::KvCache const& self) + { + auto const localityDomainId = self.localityDomainId(); + return localityDomainId ? std::optional{localityDomainId->value()} : std::nullopt; + }) .def( "plan_committed_block_drop", [](kv::KvCache& self) { return self.planCommittedBlockDrop(); }, nb::call_guard()) @@ -2544,6 +2579,25 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) return std::move(utilization.raw()); }, nb::arg("manager"), nb::arg("cache_level") = kv::kHotLevel.value(), nb::call_guard()); + mIntrospection.def( + "locality_domain_free_slots", + [](kv::KvCacheManager& manager) + { + auto const& storage = manager.storage(); + std::vector freeSlots; + for (kv::LocalityDomainId localityDomainId{0}; localityDomainId < storage.numLocalityDomains(); + ++localityDomainId) + { + kv::SlotCount free = 0; + for (kv::PoolGroupIndex pgIdx{0}; pgIdx < storage.numPoolGroups(kv::kHotLevel); ++pgIdx) + { + free += storage.getStatistics(kv::kHotLevel, pgIdx, localityDomainId).free; + } + freeSlots.push_back(free); + } + return freeSlots; + }, + nb::arg("manager"), nb::call_guard()); mIntrospection.def( "grains_for_slots", [](kv::SlotCount numSlots, std::vector const& slotSizeList, size_t granularity) @@ -2898,6 +2952,9 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) .def("is_stats_excluded", &kv::KvCacheManager::isStatsExcluded, nb::arg("kv_cache_id").none(), nb::call_guard()) .def_prop_ro("tokens_per_block", &kv::KvCacheManager::tokensPerBlock) + .def_prop_ro("num_locality_domains", &kv::KvCacheManager::numLocalityDomains) + .def("get_mapped_buffer_views", &kv::KvCacheManager::getMappedBufferViews, nb::arg("layer_id"), nb::arg("role"), + nb::call_guard()) .def_prop_ro("event_manager", [](kv::KvCacheManager const& self) { return std::dynamic_pointer_cast(self.eventSink()); }) diff --git a/cpp/tensorrt_llm/runtime/CMakeLists.txt b/cpp/tensorrt_llm/runtime/CMakeLists.txt index 9ac7bac0495e..88a5c63af15c 100644 --- a/cpp/tensorrt_llm/runtime/CMakeLists.txt +++ b/cpp/tensorrt_llm/runtime/CMakeLists.txt @@ -20,7 +20,6 @@ set(SRCS utils/numpyUtils.cpp utils/runtimeUtils.cpp utils/debugUtils.cu - locality_domain/locality_domain_utils.cpp loraUtils.cpp loraModule.cpp loraCache.cpp @@ -63,6 +62,12 @@ if(NOT WIN32) target_link_libraries(tllm_runtime_topology PUBLIC NUMA::NUMA) endif() +add_library(tllm_runtime_locality_domain STATIC + locality_domain/locality_domain_utils.cpp) +add_library(tllm::runtime_locality_domain ALIAS tllm_runtime_locality_domain) +target_link_libraries(tllm_runtime_locality_domain PUBLIC tllm::common_cuda + CUDA::cuda_driver) + add_library(tllm_runtime_memory STATIC ${RUNTIME_MEMORY_SRCS}) add_library(tllm::runtime_memory ALIAS tllm_runtime_memory) target_link_libraries( @@ -87,8 +92,9 @@ if(ENABLE_MULTI_DEVICE) target_link_libraries(tllm_runtime_distributed PUBLIC MPI::MPI_C) endif() -foreach(component tllm_runtime_memory tllm_runtime_buffers - tllm_runtime_distributed tllm_runtime_topology) +foreach(component + tllm_runtime_memory tllm_runtime_buffers tllm_runtime_distributed + tllm_runtime_topology tllm_runtime_locality_domain) target_compile_features(${component} PUBLIC cxx_std_20) set_target_properties(${component} PROPERTIES POSITION_INDEPENDENT_CODE ON) endforeach() diff --git a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt index 69f9e6a6aa02..8d6a13710dfd 100644 --- a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt +++ b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt @@ -275,6 +275,19 @@ add_tllm_gtest( LABELS "unit;batch-manager;host") +add_tllm_gtest( + kvCacheManagerV2LocalityTest + OWNER + batch-manager + SOURCES + kvCacheManagerV2LocalityTest.cpp + LINK_LIBRARIES + tllm::kv_cache + $ + LABELS + "unit;batch-manager;gpu" + REQUIRES_GPU) + add_tllm_gtest( kvCacheManagerV2ReentrantSharedMutexTest OWNER diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2LocalityTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2LocalityTest.cpp new file mode 100644 index 000000000000..46b3e3d616d9 --- /dev/null +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2LocalityTest.cpp @@ -0,0 +1,536 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "kvCacheManagerV2TestUtils.h" +#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h" +#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/exceptions.h" +#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h" +#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h" +#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/localityDomain.h" +#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.h" +#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace +{ + +using namespace tensorrt_llm::batch_manager::kv_cache_manager_v2; + +std::vector freeSlotsPerDomain(StorageManager const& storage) +{ + std::vector freeSlots; + for (LocalityDomainId localityDomainId{0}; localityDomainId < storage.numLocalityDomains(); ++localityDomainId) + { + SlotCount free = 0; + for (PoolGroupIndex pgIdx{0}; pgIdx < storage.numPoolGroups(kHotLevel); ++pgIdx) + { + free += storage.getStatistics(kHotLevel, pgIdx, localityDomainId).free; + } + freeSlots.push_back(free); + } + return freeSlots; +} + +// --------------------------------------------------------------------------- +// SlotIdMap / SlotAllocator (no GPU) +// --------------------------------------------------------------------------- + +TEST(KvCacheManagerV2LocalityTest, SlotIdMapRoundTripsAndRejectsForeignChunks) +{ + SlotIdMap map{SlotCount{4}}; + map.append(2); + map.append(0); + + EXPECT_EQ(map.slotId(0), SlotId{8}); + EXPECT_EQ(map.slotId(3), SlotId{11}); + EXPECT_EQ(map.slotId(4), SlotId{0}); + EXPECT_EQ(map.localIndex(SlotId{9}), SlotCount{1}); + EXPECT_EQ(map.localIndex(SlotId{3}), SlotCount{7}); + EXPECT_FALSE(map.localIndex(SlotId{4}).has_value()); + EXPECT_FALSE(map.localIndex(SlotId{-1}).has_value()); + + EXPECT_EQ(map.pop(), 0U); + EXPECT_FALSE(map.localIndex(SlotId{0}).has_value()); +} + +TEST(KvCacheManagerV2LocalityTest, MappedAllocatorsIssueDisjointGlobalIdsAndStampTheDomain) +{ + SlotIdMap map0{SlotCount{4}}; + SlotIdMap map1{SlotCount{4}}; + map0.append(0); + map1.append(1); + SlotAllocator alloc0{SlotCount{4}, &map0, LocalityDomainId{0}}; + SlotAllocator alloc1{SlotCount{4}, &map1, LocalityDomainId{1}}; + + std::vector slots0; + std::vector slots1; + for (int i = 0; i < 4; ++i) + { + slots0.push_back(alloc0.allocate()); + slots1.push_back(alloc1.allocate()); + } + std::set ids; + for (auto const& slot : slots0) + { + EXPECT_EQ(slot.localityDomainId(), std::optional{0}); + EXPECT_LT(slot.slotId(), SlotId{4}); + ids.insert(slot.slotId()); + } + for (auto const& slot : slots1) + { + EXPECT_EQ(slot.localityDomainId(), std::optional{1}); + EXPECT_GE(slot.slotId(), SlotId{4}); + ids.insert(slot.slotId()); + } + EXPECT_EQ(ids.size(), 8U); + EXPECT_EQ(alloc0.numFreeSlots(), 0); + + // A handle carrying the other domain's id is foreign to this allocator; the real slot stays + // with alloc1 and is released below. + { + Slot foreign; + foreign.setSlotId(slots1[0].slotId()); + foreign.setLocalityDomainId(LocalityDomainId{1}); + EXPECT_THROW(alloc0.release(std::move(foreign)), LogicError); + } + for (auto& slot : slots0) + { + Slot moved; + moved.setSlot(slot); + alloc0.release(std::move(moved)); + } + for (auto& slot : slots1) + { + Slot moved; + moved.setSlot(slot); + alloc1.release(std::move(moved)); + } + EXPECT_EQ(alloc0.numFreeSlots(), 4); + EXPECT_EQ(alloc1.numFreeSlots(), 4); + alloc0.synchronize(); + alloc1.synchronize(); + alloc0.prepareForShrink(SlotCount{0}); + EXPECT_TRUE(alloc0.finishShrink()); + alloc1.prepareForShrink(SlotCount{0}); + EXPECT_TRUE(alloc1.finishShrink()); +} + +TEST(KvCacheManagerV2LocalityTest, MappedAllocatorShrinksByLocalIndex) +{ + // Chunk 5 backs local indices [0, 4): global ids 20..23. Shrinking to 2 local slots must treat + // global ids 22 and 23 as overflow, not ids >= 2. + SlotIdMap map{SlotCount{4}}; + map.append(5); + SlotAllocator alloc{SlotCount{4}, &map, LocalityDomainId{0}}; + auto slots = std::vector{}; + for (int i = 0; i < 4; ++i) + { + slots.push_back(alloc.allocate()); + } + for (auto& slot : slots) + { + Slot moved; + moved.setSlot(slot); + alloc.release(std::move(moved)); + } + alloc.synchronize(); + alloc.prepareForShrink(SlotCount{2}); + EXPECT_EQ(alloc.numOverflowSlots(), 2); + EXPECT_TRUE(alloc.finishShrink()); + EXPECT_EQ(alloc.numSlots(), 2); + auto kept = alloc.allocate(); + EXPECT_LT(kept.slotId(), SlotId{22}); + Slot moved; + moved.setSlot(kept); + alloc.release(std::move(moved)); + alloc.synchronize(); + alloc.prepareForShrink(SlotCount{0}); + EXPECT_TRUE(alloc.finishShrink()); +} + +// --------------------------------------------------------------------------- +// Manager under TRT_LLM_MOCK_LOCALIZATION_SUPPORT=1 (GPU, ordinary device memory) +// --------------------------------------------------------------------------- + +class KvCacheManagerV2MockLocalityTest : public ::testing::Test +{ +protected: + void SetUp() override + { + ASSERT_EQ(cudaSetDevice(0), cudaSuccess); + ASSERT_EQ(cudaStreamCreate(&mStream), cudaSuccess); + mSavedMode = std::getenv("TRT_LLM_MOCK_LOCALIZATION_SUPPORT") == nullptr + ? std::nullopt + : std::optional{std::getenv("TRT_LLM_MOCK_LOCALIZATION_SUPPORT")}; + setenv("TRT_LLM_MOCK_LOCALIZATION_SUPPORT", "1", 1); + } + + void TearDown() override + { + if (mSavedMode.has_value()) + { + setenv("TRT_LLM_MOCK_LOCALIZATION_SUPPORT", mSavedMode->c_str(), 1); + } + else + { + unsetenv("TRT_LLM_MOCK_LOCALIZATION_SUPPORT"); + } + cudaStreamDestroy(mStream); + } + + static KVCacheManagerConfig makeLocalizedConfig(bool enableLocalityDomains = true) + { + KVCacheManagerConfig config; + config.tokensPerBlock = 32; + GpuCacheTierConfig gpu{256ULL << 20}; + gpu.enableLocalityDomains = enableLocalityDomains; + config.cacheTiers.emplace_back(gpu); + for (int layerId = 0; layerId < 2; ++layerId) + { + AttentionLayerConfig layer; + layer.layerId = layerId; + layer.buffers.push_back(BufferConfig{"key", 8192, std::nullopt}); + layer.buffers.push_back(BufferConfig{"value", 8192, std::nullopt}); + config.layers.emplace_back(std::move(layer)); + } + return config; + } + + static ReuseScope scopeFor(std::optional localityDomainId) + { + return ReuseScope{std::nullopt, std::nullopt, std::optional(localityDomainId)}; + } + + std::vector tokens(int count) const + { + std::vector result; + result.reserve(count); + for (int i = 0; i < count; ++i) + { + result.emplace_back(TokenId{i + 1}); + } + return result; + } + + void commitAndClose(KvCache& cache, std::vector const& allTokens, int capacity) + { + ASSERT_TRUE(cache.resume(reinterpret_cast(mStream))); + ASSERT_TRUE(cache.resize(capacity)); + int const committed = cache.numCommittedTokens(); + if (committed < static_cast(allTokens.size())) + { + cache.commit(TokenSpan{allTokens.data() + committed, static_cast(allTokens.size()) - committed}); + } + cache.stopCommitting(); + cache.close(); + } + + cudaStream_t mStream{}; + std::optional mSavedMode; +}; + +TEST_F(KvCacheManagerV2MockLocalityTest, MockReportsTwoDomainsOnlyWhenEnabled) +{ + EXPECT_EQ(std::make_shared(makeLocalizedConfig(false))->numLocalityDomains(), 1); + EXPECT_EQ(std::make_shared(makeLocalizedConfig(true))->numLocalityDomains(), 2); + EXPECT_EQ(resolveNumLocalityDomains(false), 1); +} + +TEST_F(KvCacheManagerV2MockLocalityTest, FreeSlotsAreTrackedPerDomain) +{ + auto manager = std::make_shared(makeLocalizedConfig()); + auto const before = freeSlotsPerDomain(manager->storage()); + ASSERT_EQ(before.size(), 2U); + EXPECT_GT(before[0], 0); + EXPECT_EQ(before[0], before[1]); + + auto cache = manager->createKvCache(scopeFor(1)); + EXPECT_EQ(cache->localityDomainId(), std::optional{1}); + ASSERT_TRUE(cache->resume(reinterpret_cast(mStream))); + ASSERT_TRUE(cache->resize(64)); + auto const during = freeSlotsPerDomain(manager->storage()); + EXPECT_EQ(during[0], before[0]); + EXPECT_LT(during[1], before[1]); + + cache->close(); + EXPECT_EQ(freeSlotsPerDomain(manager->storage()), before); + + EXPECT_THROW(manager->createKvCache(scopeFor(2)), tensorrt_llm::common::TllmException); + EXPECT_THROW(manager->createKvCache(scopeFor(std::nullopt)), tensorrt_llm::common::TllmException); + auto plain = std::make_shared(makeLocalizedConfig(false)); + EXPECT_THROW(plain->createKvCache(scopeFor(0)), tensorrt_llm::common::TllmException); +} + +TEST_F(KvCacheManagerV2MockLocalityTest, PageIndexUpperBoundCoversBothDomains) +{ + auto manager = std::make_shared(makeLocalizedConfig()); + auto cache0 = manager->createKvCache(scopeFor(0)); + auto cache1 = manager->createKvCache(scopeFor(1)); + ASSERT_TRUE(cache0->resume(reinterpret_cast(mStream))); + ASSERT_TRUE(cache1->resume(reinterpret_cast(mStream))); + ASSERT_TRUE(cache0->resize(128)); + ASSERT_TRUE(cache1->resize(128)); + + auto const upperBound = static_cast(manager->getPageIndexUpperBound(0, "key")); + int maxIndex = -1; + for (auto const* cache : {cache0.get(), cache1.get()}) + { + for (int index : cache->getBasePageIndices(LayerGroupId{0})) + { + if (index != kBadPageIndex.value()) + { + maxIndex = std::max(maxIndex, index); + } + } + } + EXPECT_GE(maxIndex, 0); + EXPECT_LT(maxIndex, upperBound); + + // Both domains map into one address range; the mapped extents cover every page. + size_t mappedPages = 0; + for (auto const& [address, numPages] : manager->getMappedBufferViews(0, "key")) + { + EXPECT_NE(address, 0U); + mappedPages += numPages; + } + EXPECT_GT(mappedPages, 0U); + cache0->close(); + cache1->close(); +} + +TEST_F(KvCacheManagerV2MockLocalityTest, PrefixReuseStaysWithinADomain) +{ + auto manager = std::make_shared(makeLocalizedConfig()); + auto const all = tokens(64); + TokenSpan const prefix{all.data(), static_cast(all.size()) - 1}; + + auto first = manager->createKvCache(scopeFor(0), prefix); + commitAndClose(*first, all, 128); + + auto sameDomain = manager->createKvCache(scopeFor(0), prefix); + EXPECT_GT(sameDomain->numCommittedTokens(), 0); + commitAndClose(*sameDomain, all, 128); + + auto otherDomain = manager->createKvCache(scopeFor(1), prefix); + EXPECT_EQ(otherDomain->numCommittedTokens(), 0); + commitAndClose(*otherDomain, all, 128); + + auto otherDomainAgain = manager->createKvCache(scopeFor(1), prefix); + EXPECT_GT(otherDomainAgain->numCommittedTokens(), 0); + otherDomainAgain->close(); +} + +TEST_F(KvCacheManagerV2MockLocalityTest, ResizeKeepsDomainsIndependent) +{ + auto manager = std::make_shared(makeLocalizedConfig()); + auto const all = tokens(64); + TokenSpan const prefix{all.data(), static_cast(all.size()) - 1}; + for (int domain : {0, 1}) + { + auto cache = manager->createKvCache(scopeFor(domain), prefix); + commitAndClose(*cache, all, static_cast(all.size())); + } + + auto& storage = manager->storage(); + PoolGroupIndex const pg{0}; + SlotCount const before0 = storage.numSlots(pg); + storage.shrinkPoolGroup(kHotLevel, pg, SlotCount{1}, {}, LocalityDomainId{1}); + EXPECT_EQ(storage.numSlots(pg), before0); + EXPECT_EQ(freeSlotsPerDomain(storage)[1], 0); + + // Domain 0's cached prefix survived domain 1's shrink. + auto cache = manager->createKvCache(scopeFor(0), prefix); + EXPECT_GT(cache->numCommittedTokens(), 0); + cache->close(); + + storage.expandPoolGroup(kHotLevel, pg, before0, LocalityDomainId{1}); + EXPECT_GT(freeSlotsPerDomain(storage)[1], 0); +} + +// A slot size with only a 16-byte power-of-two factor needs 2 MiB / 16 slots per chunk, about 16 GiB, so the +// manager keeps one domain instead of stranding that much memory per domain. +TEST_F(KvCacheManagerV2MockLocalityTest, OversizedChunkFallsBackToOneDomain) +{ + auto makeConfig = [](size_t slotSize) + { + KVCacheManagerConfig config; + config.tokensPerBlock = 32; + GpuCacheTierConfig gpu{256ULL << 20}; + gpu.enableLocalityDomains = true; + config.cacheTiers.emplace_back(gpu); + AttentionLayerConfig layer; + layer.layerId = 0; + layer.buffers.push_back(BufferConfig{"key", slotSize, std::nullopt}); + config.layers.emplace_back(std::move(layer)); + return config; + }; + EXPECT_EQ(std::make_shared(makeConfig(16384))->numLocalityDomains(), 2); + EXPECT_EQ(std::make_shared(makeConfig(8191 * 16))->numLocalityDomains(), 1); +} + +// One 16 KiB-slot pool, 2 MiB physical grains, two domains: the per-domain share is the raw +// quota halved and rounded down to a grain, both at construction and on resize. +TEST_F(KvCacheManagerV2MockLocalityTest, QuotaIsSharedPerDomainConsistentlyOnInitAndResize) +{ + auto makeConfig = [](size_t quota) + { + KVCacheManagerConfig config; + config.tokensPerBlock = 32; + GpuCacheTierConfig gpu{quota}; + gpu.enableLocalityDomains = true; + config.cacheTiers.emplace_back(gpu); + AttentionLayerConfig layer; + layer.layerId = 0; + layer.buffers.push_back(BufferConfig{"key", 16384, std::nullopt}); + config.layers.emplace_back(std::move(layer)); + return config; + }; + constexpr size_t kMiB = 1ULL << 20; + constexpr SlotCount kSlotsPerGrain = 128; + PoolGroupIndex const pg{0}; + + for (auto const& [quotaMiB, grainsPerDomain] : std::vector>{{6, 1}, {7, 1}, {8, 2}}) + { + auto manager = std::make_shared(makeConfig(quotaMiB * kMiB)); + auto& storage = manager->storage(); + EXPECT_EQ(storage.numSlots(pg), grainsPerDomain * kSlotsPerGrain) << quotaMiB << " MiB at init"; + for (auto const& [resizeMiB, resizeGrains] : std::vector>{{8, 2}, {7, 1}, {6, 1}}) + { + storage.adjustCacheLevel(kHotLevel, resizeMiB * kMiB, storage.getRatioList(kHotLevel), nullptr); + EXPECT_EQ(storage.numSlots(pg), resizeGrains * kSlotsPerGrain) << resizeMiB << " MiB on resize"; + auto const free = freeSlotsPerDomain(storage); + EXPECT_EQ(free[0], free[1]); + EXPECT_EQ(free[0], resizeGrains * kSlotsPerGrain); + } + } +} + +TEST_F(KvCacheManagerV2MockLocalityTest, ManagerResizeDecidesPerDomain) +{ + constexpr size_t kFullQuota = 256ULL << 20; + constexpr size_t kMidQuota = 192ULL << 20; + auto manager = std::make_shared(makeLocalizedConfig()); + auto& storage = manager->storage(); + PoolGroupIndex const pg{0}; + auto const ratio = storage.getRatioList(kHotLevel); + SlotCount const full = storage.numSlots(pg); + ASSERT_TRUE(manager->resize(kHotLevel, kMidQuota)); + SlotCount const mid = storage.numSlots(pg); + ASSERT_LT(mid, full); + ASSERT_TRUE(manager->resize(kHotLevel, kFullQuota)); + ASSERT_EQ(storage.numSlots(pg), full); + + // Domain 0 above the target and domain 1 below it: one shrinks while the other grows. + storage.shrinkPoolGroup(kHotLevel, pg, mid / 2, {}, LocalityDomainId{1}); + ASSERT_EQ(freeSlotsPerDomain(storage)[1], mid / 2); + storage.adjustCacheLevel(kHotLevel, kMidQuota, ratio, nullptr); + EXPECT_EQ(freeSlotsPerDomain(storage), (std::vector{mid, mid})); + + // Domain 1 already at the target: only domain 0 moves. + storage.shrinkPoolGroup(kHotLevel, pg, mid / 2, {}, LocalityDomainId{0}); + storage.adjustCacheLevel(kHotLevel, kMidQuota, ratio, nullptr); + EXPECT_EQ(freeSlotsPerDomain(storage), (std::vector{mid, mid})); + + EXPECT_THROW(storage.expandPoolGroup(kHotLevel, pg, mid, LocalityDomainId{0}), std::exception); + EXPECT_THROW(storage.shrinkPoolGroup(kHotLevel, pg, mid, {}, LocalityDomainId{1}), std::exception); + EXPECT_EQ(freeSlotsPerDomain(storage), (std::vector{mid, mid})); +} + +// A 3 MiB slot does not divide the 2 MiB grain, so each domain's backing rounds up on its own. +// Re-applying the reported quota must not lose slots. +TEST_F(KvCacheManagerV2MockLocalityTest, QuotaRoundTripKeepsCoarseSlots) +{ + constexpr size_t kMiB = 1ULL << 20; + KVCacheManagerConfig config; + config.tokensPerBlock = 32; + GpuCacheTierConfig gpu{32 * kMiB}; + gpu.enableLocalityDomains = true; + config.cacheTiers.emplace_back(gpu); + AttentionLayerConfig layer; + layer.layerId = 0; + layer.buffers.push_back(BufferConfig{"key", 3 * kMiB, std::nullopt}); + config.layers.emplace_back(std::move(layer)); + + auto manager = std::make_shared(config); + auto& storage = manager->storage(); + PoolGroupIndex const pg{0}; + SlotCount const initial = storage.numSlots(pg); + ASSERT_GT(initial, 0); + size_t const perDomainBytes = (static_cast(initial) * 3 * kMiB + 2 * kMiB - 1) / (2 * kMiB) * (2 * kMiB); + EXPECT_EQ(manager->getQuota(kHotLevel), 2 * perDomainBytes); + for (int i = 0; i < 2; ++i) + { + manager->adjust(); + EXPECT_EQ(storage.numSlots(pg), initial); + EXPECT_EQ(freeSlotsPerDomain(storage), (std::vector{initial, initial})); + EXPECT_EQ(manager->getQuota(kHotLevel), 2 * perDomainBytes); + } +} + +TEST_F(KvCacheManagerV2MockLocalityTest, SecondaryGpuTierIsRejectedWhenLocalized) +{ + auto config = makeLocalizedConfig(); + config.cacheTiers.emplace_back(GpuCacheTierConfig{64ULL << 20}); + try + { + KvCacheManager manager(config); + FAIL() << "a localized manager accepted a secondary GPU tier"; + } + catch (std::invalid_argument const& e) + { + EXPECT_NE(std::string(e.what()).find("Locality domains"), std::string::npos) << e.what(); + } +} + +TEST_F(KvCacheManagerV2MockLocalityTest, UtilizationIsReportedPerDomain) +{ + auto manager = std::make_shared(makeLocalizedConfig()); + auto& storage = manager->storage(); + auto const capacity = storage.numSlots(PoolGroupIndex{0}); + auto cache = manager->createKvCache(scopeFor(1)); + ASSERT_TRUE(cache->resume(reinterpret_cast(mStream))); + ASSERT_TRUE(cache->resize(static_cast(capacity) / 2 * 32)); + + EXPECT_FLOAT_EQ(storage.getUtilization(kHotLevel, LocalityDomainId{0}).at(PoolGroupIndex{0}), 0.f); + EXPECT_GT(storage.getUtilization(kHotLevel, LocalityDomainId{1}).at(PoolGroupIndex{0}), 0.4f); + EXPECT_LT(storage.getUtilization(kHotLevel).at(PoolGroupIndex{0}), 0.35f); + cache->close(); +} + +TEST_F(KvCacheManagerV2MockLocalityTest, NonLocalizedScopesHashAsBefore) +{ + // The third ReuseScope field must not change the key of scopes that do not set it. + ReuseScope const plain{7, "11"}; + ReuseScope const explicitNone{7, "11", std::nullopt}; + EXPECT_EQ(Hasher(plain).digest(), Hasher(explicitNone).digest()); + EXPECT_NE(Hasher(plain).digest(), Hasher(ReuseScope{7, "11", LocalityDomainId{0}}).digest()); + EXPECT_NE(Hasher(ReuseScope{7, "11", LocalityDomainId{0}}).digest(), + Hasher(ReuseScope{7, "11", LocalityDomainId{1}}).digest()); +} + +} // namespace diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2TypedIndexTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2TypedIndexTest.cpp index 6a4833aff659..ffb65a7cdc0d 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2TypedIndexTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2TypedIndexTest.cpp @@ -27,6 +27,7 @@ #include #include +#include #include #include @@ -245,13 +246,14 @@ TEST(KvCacheManagerV2TypedIndexTest, AllocationRequestsUsePlainSlotCount) static_assert(!std::is_invocable::value, "SlotAllocator::allocateMultiple must not take a slot id"); - using StorageManagerNewGpuSlots = TypedVec> (StorageManager::*)( - TypedVec const&, MigrationRecorder const&, DropRecorder const&); + using StorageManagerNewGpuSlots + = TypedVec> (StorageManager::*)(TypedVec const&, + MigrationRecorder const&, DropRecorder const&, std::optional); static_assert(std::is_same::value, "StorageManager::newGpuSlots must take plain slot counts"); - using StorageManagerNewSlotsForPoolGroup = std::vector (StorageManager::*)( - CacheLevel, PoolGroupIndex, SlotCount, MigrationRecorder const&, DropRecorder const&); + using StorageManagerNewSlotsForPoolGroup = std::vector (StorageManager::*)(CacheLevel, PoolGroupIndex, + SlotCount, MigrationRecorder const&, DropRecorder const&, std::optional); static_assert( std::is_same::value, "StorageManager::newSlotsForPoolGroup must take a plain slot count"); diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index 2f1644a37c31..4f1e58139d1b 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -134,11 +134,14 @@ class ReuseScope: self, lora_id: int | None = None, salt: str | None = None, + locality_domain_id: int | None = None, ) -> None: ... @property def lora_id(self) -> int | None: ... @property def salt(self) -> str | None: ... + @property + def locality_domain_id(self) -> int | None: ... def __len__(self) -> int: ... def __getitem__(self, index: int) -> int | str | None: ... def __iter__(self) -> Iterator[int | str | None]: ... @@ -236,6 +239,7 @@ class CacheTierConfig(Protocol): @dataclass(slots=True) class GpuCacheTierConfig: quota: int + enable_locality_domains: bool = False @property def tier(self) -> CacheTier: ... def validate(self) -> None: ... @@ -669,6 +673,8 @@ class _KVCache: def committed_tokens(self) -> list[TokenIdExt]: ... @property def reuse_scope(self) -> ReuseScope: ... + @property + def locality_domain_id(self) -> int | None: ... def plan_committed_block_drop(self) -> PlannedDropHandle | None: ... def stop_committing(self) -> None: ... def suspend(self) -> None: ... @@ -725,6 +731,7 @@ class PoolDesc: pool_index: PoolIndex base_address: MemAddress slot_bytes: int + mapped_ranges: list[tuple[int, int]] class BufferId(NamedTuple): layer_id: LayerId @@ -914,6 +921,11 @@ class KVCacheManager: @property def tokens_per_block(self) -> int: ... @property + def num_locality_domains(self) -> int: ... + def get_mapped_buffer_views( + self, layer_id: LayerId, role: DataRole + ) -> list[tuple[MemAddress, int]]: ... + @property def event_manager(self) -> Any | None: ... @property def init_config(self) -> KVCacheManagerConfig: ... diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py index 4374facb4052..7ba3c2a5ed22 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py @@ -246,6 +246,11 @@ def storage_utilization(manager: Any, cache_level: int = 0) -> list[float]: return list(_cpp().storage_utilization(manager, cache_level)) +def locality_domain_free_slots(manager: Any) -> list[int]: + """Return the free GPU slot count per locality domain, summed over pool groups.""" + return list(_cpp().locality_domain_free_slots(manager)) + + def grains_for_slots(num_slots: int, slot_size_list: list[int], granularity: int) -> int: """Return the grain count required for a pool group slot count.""" return int(_cpp().grains_for_slots(num_slots, slot_size_list, granularity)) diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_localization.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_localization.py new file mode 100644 index 000000000000..dfe6c021c46e --- /dev/null +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_localization.py @@ -0,0 +1,302 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Locality domain behavior of KVCacheManager and its KV caches. + +Localized tests opt in with GpuCacheTierConfig(enable_locality_domains=True) and set +TRT_LLM_MOCK_LOCALIZATION_SUPPORT=1, so the manager reports two domains backed by +ordinary device memory. Physical placement is therefore not verified here; the +metadata, addressing, reuse isolation and per-domain capacity are. +""" + +import gc +import os +import sys +import unittest +from contextlib import contextmanager +from importlib.util import find_spec +from typing import TYPE_CHECKING, cast + +if not TYPE_CHECKING and find_spec("kv_cache_manager_v2") is not None: + from kv_cache_manager_v2 import ( + BAD_PAGE_INDEX, + AttentionLayerConfig, + BufferConfig, + CudaStream, + DataRole, + GpuCacheTierConfig, + KVCacheManager, + KVCacheManagerConfig, + LayerId, + ReuseScope, + _introspection, + _KVCache, + ) +else: + from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + BAD_PAGE_INDEX, + AttentionLayerConfig, + BufferConfig, + CudaStream, + DataRole, + GpuCacheTierConfig, + KVCacheManager, + KVCacheManagerConfig, + LayerId, + ReuseScope, + _introspection, + _KVCache, + ) + +_TEST_DIR = os.path.dirname(os.path.abspath(__file__)) +_ADDED_TEST_DIR = _TEST_DIR not in sys.path +if _ADDED_TEST_DIR: + sys.path.insert(0, _TEST_DIR) +try: + from cuda_test_utils import TemporaryCudaStream, init_cuda_once # noqa: E402 +finally: + if _ADDED_TEST_DIR: + sys.path.remove(_TEST_DIR) + +ROLE_KEY = DataRole("key") +ROLE_VALUE = DataRole("value") + +_256MB = 256 * (1 << 20) +_TOKENS_PER_BLOCK = 32 +_KV_BUF_SIZE = 8192 + + +@contextmanager +def _override_localization_mode(mode: str): + env_backup = os.environ.get("TRT_LLM_MOCK_LOCALIZATION_SUPPORT") + try: + os.environ["TRT_LLM_MOCK_LOCALIZATION_SUPPORT"] = mode + yield + finally: + if env_backup is None: + os.environ.pop("TRT_LLM_MOCK_LOCALIZATION_SUPPORT", None) + else: + os.environ["TRT_LLM_MOCK_LOCALIZATION_SUPPORT"] = env_backup + + +def _make_config( + num_layers: int = 2, + gpu_quota: int = _256MB, + tokens_per_block: int = _TOKENS_PER_BLOCK, + kv_buf_size: int = _KV_BUF_SIZE, + enable_locality_domains: bool = False, +) -> KVCacheManagerConfig: + return KVCacheManagerConfig( + tokens_per_block=tokens_per_block, + cache_tiers=[ + GpuCacheTierConfig(quota=gpu_quota, enable_locality_domains=enable_locality_domains) + ], + layers=[ + AttentionLayerConfig( + layer_id=LayerId(i), + buffers=[ + BufferConfig(role=ROLE_KEY, size=kv_buf_size), + BufferConfig(role=ROLE_VALUE, size=kv_buf_size), + ], + sliding_window_size=None, + ) + for i in range(num_layers) + ], + ) + + +def _make_localized_manager(**kwargs) -> KVCacheManager: + kwargs.setdefault("enable_locality_domains", True) + config = _make_config(**kwargs) + with _override_localization_mode("1"): + return KVCacheManager(config) + + +def _make_non_localized_manager(**kwargs) -> KVCacheManager: + config = _make_config(**kwargs) + with _override_localization_mode("0"): + return KVCacheManager(config) + + +def _scope(locality_domain_id: int | None) -> ReuseScope: + return ReuseScope(locality_domain_id=locality_domain_id) + + +def _valid_page_indices(kv_cache: _KVCache, layer_group_id: int = 0) -> list[int]: + return [idx for idx in kv_cache.get_base_page_indices(layer_group_id) if idx != BAD_PAGE_INDEX] + + +class _ManagerTestCase(unittest.TestCase): + manager: KVCacheManager | None + + def setUp(self) -> None: + init_cuda_once() + gc.collect() + gc.disable() + self.manager = None + + def tearDown(self) -> None: + gc.enable() + if self.manager is not None: + self.manager.shutdown() + self.manager = None + + def _commit_and_close(self, kv_cache: _KVCache, tokens: list[int], capacity: int) -> None: + with TemporaryCudaStream([]) as s: + stream = cast(CudaStream, s.handle) + self.assertTrue(kv_cache.resume(stream)) + self.assertTrue(kv_cache.resize(capacity)) + uncommitted = tokens[kv_cache.num_committed_tokens :] + if uncommitted: + kv_cache.commit(uncommitted) + kv_cache.stop_committing() + kv_cache.close() + + +class TestLocalityDomainId(_ManagerTestCase): + def test_stored_on_creation(self) -> None: + self.manager = _make_localized_manager() + for domain in (0, 1): + kv_cache = self.manager.create_kv_cache(_scope(domain)) + self.assertEqual(kv_cache.locality_domain_id, domain) + self.assertEqual(kv_cache.reuse_scope.locality_domain_id, domain) + kv_cache.close() + + def test_none_when_not_localized(self) -> None: + self.manager = _make_non_localized_manager() + kv_cache = self.manager.create_kv_cache() + self.assertIsNone(kv_cache.locality_domain_id) + self.assertEqual(len(kv_cache.reuse_scope), 2) + kv_cache.close() + + def test_out_of_range_rejected(self) -> None: + self.manager = _make_localized_manager() + with self.assertRaises(Exception): + self.manager.create_kv_cache(_scope(2)) + self.manager = _make_non_localized_manager() + with self.assertRaises(Exception): + self.manager.create_kv_cache(_scope(0)) + + +class TestLocalizedQueries(_ManagerTestCase): + def test_num_locality_domains(self) -> None: + self.manager = _make_non_localized_manager() + self.assertEqual(self.manager.num_locality_domains, 1) + self.manager.shutdown() + self.manager = _make_localized_manager() + self.assertEqual(self.manager.num_locality_domains, 2) + + def test_free_slots_follow_the_domain(self) -> None: + self.manager = _make_localized_manager() + before = _introspection.locality_domain_free_slots(self.manager) + self.assertEqual(len(before), 2) + self.assertGreater(before[0], 0) + self.assertEqual(before[0], before[1]) + + kv_cache = self.manager.create_kv_cache(_scope(1)) + with TemporaryCudaStream([]) as s: + stream = cast(CudaStream, s.handle) + self.assertTrue(kv_cache.resume(stream)) + self.assertTrue(kv_cache.resize(64)) + during = _introspection.locality_domain_free_slots(self.manager) + self.assertEqual(during[0], before[0]) + self.assertLess(during[1], before[1]) + + kv_cache.close() + self.assertEqual(_introspection.locality_domain_free_slots(self.manager), before) + + def test_free_slots_non_localized(self) -> None: + self.manager = _make_non_localized_manager() + free_slots = _introspection.locality_domain_free_slots(self.manager) + self.assertEqual(len(free_slots), 1) + self.assertGreater(free_slots[0], 0) + + def test_page_index_upper_bound_covers_both_domains(self) -> None: + self.manager = _make_localized_manager() + caches = [self.manager.create_kv_cache(_scope(domain)) for domain in (0, 1)] + with TemporaryCudaStream([]) as s: + stream = cast(CudaStream, s.handle) + for kv_cache in caches: + self.assertTrue(kv_cache.resume(stream)) + self.assertTrue(kv_cache.resize(128)) + upper_bound = self.manager.get_page_index_upper_bound(LayerId(0), ROLE_KEY) + indices = [idx for kv_cache in caches for idx in _valid_page_indices(kv_cache)] + self.assertTrue(indices) + self.assertLess(max(indices), upper_bound) + self.assertEqual(len(set(indices)), len(indices)) + for kv_cache in caches: + kv_cache.close() + + def test_mapped_buffer_views_cover_the_pool(self) -> None: + self.manager = _make_localized_manager() + views = self.manager.get_mapped_buffer_views(LayerId(0), ROLE_KEY) + self.assertTrue(views) + for address, num_pages in views: + self.assertGreater(address, 0) + self.assertGreater(num_pages, 0) + self.manager.shutdown() + self.manager = _make_non_localized_manager() + self.assertEqual(len(self.manager.get_mapped_buffer_views(LayerId(0), ROLE_KEY)), 1) + + +class TestLocalizedReuse(_ManagerTestCase): + def test_same_domain_reuses_blocks(self) -> None: + self.manager = _make_localized_manager() + tokens = list(range(64)) + kv_a = self.manager.create_kv_cache(_scope(0), tokens[:-1]) + self._commit_and_close(kv_a, tokens, 128) + kv_b = self.manager.create_kv_cache(_scope(0), tokens[:-1]) + self.assertGreater(kv_b.num_committed_tokens, 0) + self._commit_and_close(kv_b, tokens, 128) + + def test_other_domain_does_not_reuse_blocks(self) -> None: + self.manager = _make_localized_manager() + tokens = list(range(64)) + kv_a = self.manager.create_kv_cache(_scope(0), tokens[:-1]) + self._commit_and_close(kv_a, tokens, 128) + kv_b = self.manager.create_kv_cache(_scope(1), tokens[:-1]) + self.assertEqual(kv_b.num_committed_tokens, 0) + self._commit_and_close(kv_b, tokens, 128) + kv_c = self.manager.create_kv_cache(_scope(1), tokens[:-1]) + self.assertGreater(kv_c.num_committed_tokens, 0) + kv_c.close() + + def test_non_localized_reuse_unaffected(self) -> None: + self.manager = _make_non_localized_manager() + tokens = list(range(64)) + kv_a = self.manager.create_kv_cache(None, tokens[:-1]) + self._commit_and_close(kv_a, tokens, 128) + kv_b = self.manager.create_kv_cache(None, tokens[:-1]) + self.assertGreater(kv_b.num_committed_tokens, 0) + self._commit_and_close(kv_b, tokens, 128) + + +class TestLocalizedSuspendResume(_ManagerTestCase): + def test_suspend_resume_preserves_domain_and_pages(self) -> None: + self.manager = _make_localized_manager() + kv_cache = self.manager.create_kv_cache(_scope(1)) + with TemporaryCudaStream([]) as s: + stream = cast(CudaStream, s.handle) + self.assertTrue(kv_cache.resume(stream)) + self.assertTrue(kv_cache.resize(128)) + before = _valid_page_indices(kv_cache) + kv_cache.suspend() + self.assertTrue(kv_cache.resume(stream)) + self.assertEqual(kv_cache.locality_domain_id, 1) + self.assertEqual(_valid_page_indices(kv_cache), before) + kv_cache.close() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_localization_wrapper.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_localization_wrapper.py new file mode 100644 index 000000000000..f3aff5db2a6f --- /dev/null +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_localization_wrapper.py @@ -0,0 +1,176 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""KVCacheManagerV2 wrapper paths that build a ReuseScope for a request. + +Every probe, prefetch and mirror creation must carry the request's locality +domain, or the native manager rejects the scope (localized) and the lookup +misses the domain-keyed blocks. The wrapper is driven as a bare instance over a +real native manager in mock localization mode, with no model. +""" + +import gc +import os +import sys +import unittest +from contextlib import contextmanager +from types import SimpleNamespace +from typing import cast + +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + AttentionLayerConfig, + BufferConfig, + CudaStream, + DataRole, + GpuCacheTierConfig, + KVCacheManager, + KVCacheManagerConfig, + LayerId, + ReuseScope, +) + +_TEST_DIR = os.path.dirname(os.path.abspath(__file__)) +_ADDED_TEST_DIR = _TEST_DIR not in sys.path +if _ADDED_TEST_DIR: + sys.path.insert(0, _TEST_DIR) +try: + from cuda_test_utils import TemporaryCudaStream, init_cuda_once # noqa: E402 +finally: + if _ADDED_TEST_DIR: + sys.path.remove(_TEST_DIR) + + +@contextmanager +def _mock_localization(): + env_backup = os.environ.get("TRT_LLM_MOCK_LOCALIZATION_SUPPORT") + os.environ["TRT_LLM_MOCK_LOCALIZATION_SUPPORT"] = "1" + try: + yield + finally: + if env_backup is None: + os.environ.pop("TRT_LLM_MOCK_LOCALIZATION_SUPPORT", None) + else: + os.environ["TRT_LLM_MOCK_LOCALIZATION_SUPPORT"] = env_backup + + +def _make_impl() -> KVCacheManager: + config = KVCacheManagerConfig( + tokens_per_block=32, + cache_tiers=[GpuCacheTierConfig(quota=256 << 20, enable_locality_domains=True)], + layers=[ + AttentionLayerConfig( + layer_id=LayerId(0), + buffers=[ + BufferConfig(role=DataRole("key"), size=8192), + BufferConfig(role=DataRole("value"), size=8192), + ], + sliding_window_size=None, + ) + ], + ) + with _mock_localization(): + return KVCacheManager(config) + + +class _RecordingImpl: + """Delegates to the native manager and records every ReuseScope it sees.""" + + def __init__(self, impl): + self._impl = impl + self.scopes = [] + + def __getattr__(self, name): + return getattr(self._impl, name) + + def create_kv_cache(self, scope, *args, **kwargs): + self.scopes.append(scope) + return self._impl.create_kv_cache(scope, *args, **kwargs) + + def probe_reuse(self, scope, tokens): + self.scopes.append(scope) + return self._impl.probe_reuse(scope, tokens) + + def probe_first_new_block_key(self, scope, tokens): + self.scopes.append(scope) + return self._impl.probe_first_new_block_key(scope, tokens) + + def prefetch_reuse(self, scope, tokens, cache_level): + # The fixture has no tier to prefetch from; only the scope matters here. + self.scopes.append(scope) + return False + + +def _request(request_id: int, domain: int): + return SimpleNamespace( + py_request_id=request_id, py_locality_domain_id=domain, lora_task_id=None, cache_salt=None + ) + + +def _bare_wrapper(impl) -> KVCacheManagerV2: + wrapper = cast(KVCacheManagerV2, object.__new__(KVCacheManagerV2)) + wrapper.impl = impl + wrapper.enable_block_reuse = True + wrapper.can_evict = True + wrapper.kv_cache_map = {} + return wrapper + + +class TestWrapperScopesCarryTheDomain(unittest.TestCase): + def setUp(self) -> None: + init_cuda_once() + gc.collect() + gc.disable() + self.impl = _RecordingImpl(_make_impl()) + self.wrapper = _bare_wrapper(self.impl) + self.tokens = list(range(64)) + self.wrapper._context_reuse_tokens = lambda req, reuse_limit=None: self.tokens[:-1] + # Commit a prefix in domain 0 only. + with TemporaryCudaStream([]) as s: + stream = cast(CudaStream, s.handle) + native = self.impl._impl.create_kv_cache( + ReuseScope(locality_domain_id=0), self.tokens[:-1] + ) + self.assertTrue(native.resume(stream)) + self.assertTrue(native.resize(128)) + native.commit(self.tokens[native.num_committed_tokens :]) + native.stop_committing() + native.close() + self.impl.scopes.clear() + + def tearDown(self) -> None: + gc.enable() + self.impl._impl.shutdown() + + def test_probe_context_reuse_sees_only_its_domain(self) -> None: + self.assertGreater(self.wrapper.probe_context_reuse(_request(1, 0)), 0) + self.assertEqual(self.wrapper.probe_context_reuse(_request(2, 1)), 0) + self.assertEqual([s.locality_domain_id for s in self.impl.scopes], [0, 1]) + + def test_probe_prefix_match_length_takes_the_best_domain(self) -> None: + match = self.wrapper.probe_prefix_match_length(self.tokens[:-1]) + self.assertGreater(match, 0) + self.assertEqual(sorted(s.locality_domain_id for s in self.impl.scopes), [0, 1]) + + def test_probe_first_new_block_key_uses_the_domain(self) -> None: + self.wrapper.probe_first_new_block_key(_request(3, 1)) + self.assertEqual([s.locality_domain_id for s in self.impl.scopes], [1]) + + def test_prefetch_builds_a_domain_scope(self) -> None: + self.wrapper.prefetch_for_context_tokens([_request(4, 1), _request(5, 0)]) + self.assertEqual([s.locality_domain_id for s in self.impl.scopes], [1, 0]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_salting.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_salting.py index e8f1bb0c43f0..fea52483fe73 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_salting.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_salting.py @@ -88,6 +88,19 @@ def test_reuse_scope_seeds_distinct_keys(self) -> None: # Deterministic across repeated derivation. self.assertEqual(roots, [_root_key(scope) for scope in scopes]) + def test_locality_domain_is_part_of_the_scope_key(self) -> None: + # A scope without a locality domain keeps its two-field shape and key, so existing + # digests are unchanged; each domain seeds its own key. + plain = ReuseScope(lora_id=7, salt="11") + self.assertEqual(len(plain), 2) + self.assertEqual(_root_key(plain), _root_key(ReuseScope(7, "11", None))) + domain_keys = { + _root_key(ReuseScope(lora_id=7, salt="11", locality_domain_id=d)) for d in (0, 1) + } + self.assertEqual(len(domain_keys), 2) + self.assertNotIn(_root_key(plain), domain_keys) + self.assertEqual(len(ReuseScope(lora_id=7, salt="11", locality_domain_id=0)), 3) + def test_blockchain_keys_are_seeded_by_reuse_scope(self) -> None: tokens = [TokenId(1), TokenId(2), TokenId(3), TokenId(4)] scope = ReuseScope(lora_id=7, salt="11") From 022ce6137594e1d78ee7baa8358e4af0d4de15d8 Mon Sep 17 00:00:00 2001 From: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> Date: Fri, 23 Oct 2026 18:38:46 +0000 Subject: [PATCH 2/5] [None][feat] Locality domains in the Python KVCacheManagerV2 and request path enable_locality_domains builds a localized KVCacheManagerV2. Each new request is placed on a domain round-robin over arrival order, carries it through every KV cache creation, probe and prefetch, and child requests inherit it. Sparse attention models are rejected at manager creation when the flag is set, and mock localization keeps attention on one stream. Signed-off-by: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/_util.py | 10 + .../pyexecutor/executor_request_queue.py | 1 + .../kv_cache/kv_cache_manager_v2.py | 257 +++++++++++++++++- .../kv_cache/mamba_cache_manager.py | 4 + tensorrt_llm/_torch/pyexecutor/llm_request.py | 8 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 6 + .../_torch/pyexecutor/request_utils.py | 7 +- .../pyexecutor/scheduler/scheduler_v2.py | 17 +- .../kv_cache/test_kv_cache_budget_split.py | 11 +- .../kv_cache/test_kv_cache_estimation.py | 1 + .../kv_cache/test_kv_cache_v2_scheduler.py | 18 ++ .../executor/test_locality_domain_request.py | 84 ++++++ .../_torch/executor/test_py_executor.py | 2 + .../_torch/executor/test_seq_slot_sizing.py | 1 + tests/unittest/llmapi/test_llm_args.py | 17 ++ 15 files changed, 427 insertions(+), 17 deletions(-) create mode 100644 tests/unittest/_torch/executor/test_locality_domain_request.py diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 9a323116808b..dd921821406a 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2032,6 +2032,11 @@ def _create_kv_cache_manager( kv_cache_config_override: Optional[KvCacheConfig] = None, cold_page_codec_provider: Optional[object] = None, ) -> KVCacheManager: + if self._llm_args.enable_locality_domains and self._sparse_attention_config is not None: + raise ValueError( + "enable_locality_domains does not support sparse attention (algorithm " + f"{getattr(self._sparse_attention_config, 'algorithm', None)!r}); only dense " + "and MLA attention are partitioned per locality domain.") mapping = self._mapping assert model_engine.model.model_config.is_generation, "Only construct KV cache for generation models." kv_cache_config = (kv_cache_config_override if kv_cache_config_override @@ -2066,6 +2071,7 @@ def _create_kv_cache_manager( enable_kv_cache_stats=self._enable_kv_cache_stats() and not estimating_kv_cache, execution_stream=self._execution_stream, + enable_locality_domains=self._llm_args.enable_locality_domains, layer_mask=spec_dec_layer_mask, is_disagg=self._is_disagg, disable_overlap_scheduler=self._disable_overlap_scheduler, @@ -2388,6 +2394,7 @@ def _create_one_model_draft_kv_cache_manager( layer_mask=spec_dec_layer_mask, num_layers=num_draft_layers, is_disagg=self._is_disagg, + enable_locality_domains=self._llm_args.enable_locality_domains, disable_overlap_scheduler=self._disable_overlap_scheduler, cold_page_codec_provider=cold_page_codec_provider, joint_kv_cache_reuse=self._joint_kv_cache_reuse, @@ -3271,6 +3278,7 @@ def _create_kv_cache_manager( estimating_kv_cache: bool = False, enable_kv_cache_stats: bool = False, execution_stream: Optional[torch.cuda.Stream] = None, + enable_locality_domains: bool = False, # Optional overrides for one-model draft case (when model_engine is None) model_config: Optional[ModelConfig] = None, dtype: Optional[torch.dtype] = None, @@ -3427,6 +3435,8 @@ def _create_kv_cache_manager( draft_config_for_kv) manager_extra_kwargs = {} if issubclass(kv_cache_manager_cls, KVCacheManagerV2): + manager_extra_kwargs[ + "enable_locality_domains"] = enable_locality_domains manager_extra_kwargs["max_cuda_graph_batch_size"] = ( model_engine._max_cuda_graph_batch_size if model_engine is not None else max_cuda_graph_batch_size) diff --git a/tensorrt_llm/_torch/pyexecutor/executor_request_queue.py b/tensorrt_llm/_torch/pyexecutor/executor_request_queue.py index fa79c1f16108..edc641b8e32d 100644 --- a/tensorrt_llm/_torch/pyexecutor/executor_request_queue.py +++ b/tensorrt_llm/_torch/pyexecutor/executor_request_queue.py @@ -41,6 +41,7 @@ class RequestQueueItem: request: Optional[ExecutorRequest] = None _ = dataclasses.KW_ONLY child_req_ids: Optional[list] = None + locality_domain_id: Optional[int] = None is_canceled_request: bool = False # Only meaningful for control requests. True = drain active/waiting # queues before firing the action. False = fire at the next scheduler diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 806e104d148a..9b3df5cbb819 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -38,6 +38,7 @@ from tensorrt_llm._torch.disaggregation.resource.page import MapperKind, RoleLayout from tensorrt_llm._torch.distributed.communicator import Distributed, ReduceOp +from tensorrt_llm._torch.locality_domain_utils import is_locality_domain_enabled from tensorrt_llm._torch.model_config import _get_kv_cache_layer_specs from tensorrt_llm._torch.utils import helix_local_len, maybe_compile from tensorrt_llm._utils import ( @@ -1285,6 +1286,42 @@ def _swa_branch_snapshot_point( return point +def resolve_fork_join_attn( + num_locality_domains: int, + enable_locality_domains: bool, + *, + runtime_enabled: bool, + mock_localization: bool, +) -> bool: + """Whether attention forks one stream per locality domain. + + A localized KV cache needs the locality domain runtime to fork; mock + localization keeps the split cache on one stream, and a disabled runtime + with the flag on is a configuration error rather than a silent downgrade. + """ + if num_locality_domains <= 1 or not enable_locality_domains: + return False + if runtime_enabled: + # Balanced generation dispatch splits rows in half (runner.py). + if num_locality_domains != 2: + raise ValueError( + "fork-join attention supports exactly 2 locality domains; " + f"got {num_locality_domains}" + ) + return True + if mock_localization: + logger.info( + "Mock localization: KV cache split across %d domains, attention stays on one stream", + num_locality_domains, + ) + return False + raise RuntimeError( + "enable_locality_domains=True, but locality domain execution is disabled on this " + "device (DISABLE_LOCALITY_DOMAINS=1 or no runtime support), so the KV cache would " + "be localized without matching streams" + ) + + class KVCacheManagerV2(BaseResourceManager): # Sparse managers attach a metadata batch during initialization. sparse_metadata_batch: Batch | None = None @@ -1353,6 +1390,7 @@ def __init__( execution_stream: Optional[torch.cuda.Stream] = None, is_disagg: bool = False, enable_stats: bool = False, + enable_locality_domains: bool = False, num_reserved_index_slots: int = 1, disable_overlap_scheduler: bool = False, kv_events_config: Optional[KVEventsConfig] = None, @@ -1695,7 +1733,9 @@ def append_to_kv_heads_per_layer( logger.info(f"KV cache manager v2 device quota set to {quota / (1 << 30)}GiB") - cache_tiers: List[CacheTierConfig] = [GpuCacheTierConfig(quota=int(quota))] + cache_tiers: List[CacheTierConfig] = [ + GpuCacheTierConfig(quota=int(quota), enable_locality_domains=enable_locality_domains) + ] # A capacity-only connector registers nothing for migration to # invalidate, so it keeps the tier the scheduler spills to. The # rejection of non-GPU tiers in `PyExecutor` exempts it to match. @@ -1839,6 +1879,15 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: assert candidate is not None self.kv_cache_manager_py_config = config self.impl = candidate + # Cached at construction so the per-locality-domain metadata invariants + # stay stable for the manager's lifetime. + self._fork_join_attn = resolve_fork_join_attn( + self.num_locality_domains, + enable_locality_domains, + runtime_enabled=is_locality_domain_enabled(), + mock_localization=os.environ.get("TRT_LLM_MOCK_LOCALIZATION_SUPPORT") == "1", + ) + self._next_locality_domain = 0 self.can_evict = len(config.cache_tiers) > 1 if self.event_manager is not None: self.event_manager.set_layer_group_window_sizes( @@ -2122,6 +2171,16 @@ def _reserve_guard_page(self) -> None: if value is None: self._guard_page_fill = "" return + if self.num_locality_domains > 1: + # The sentinel is restored per layer buffer after warmup scrubs, and a + # localized pool has no single layer buffer to index it against. + logger.warning( + "TRTLLM_KV_GUARD_PAGE is not supported with locality domains; " + "no guard page reserved" + ) + self._guard_page_fill = "" + self._guard_page_value = None + return kv_cache = self._create_kv_cache(_GUARD_PAGE_REQUEST_ID, None, [1], is_dummy=True) if kv_cache is None: logger.warning("KVCacheManagerV2: could not reserve a guard page") @@ -3745,6 +3804,68 @@ def get_index_k_buffer( full_view = convert_to_torch_tensor(TensorWrapper(addr, torch_dtype, full_slot_shape)) return full_view[:, 0] + @property + def fork_join_attn(self) -> bool: + """Whether attention runs as a per-locality-domain fork-join.""" + return self._fork_join_attn + + @property + def num_locality_domains(self) -> int: + """Number of locality domains. Returns 1 for non-localized configurations.""" + # Test stand-ins build the manager without an impl, and a Mock impl answers + # every attribute; read anything that is not an int as non-localized. + value = getattr(getattr(self, "impl", None), "num_locality_domains", 1) + return value if isinstance(value, int) else 1 + + def pick_locality_domain(self, request_id: int) -> int | None: + """Default locality domain placement for newly-created requests. + + Round-robin over arrival order. Request ids are not used: ids that + advance in even steps (a parent and its child per request, or + round-robin disagg ids) would otherwise pin every request to one domain. + """ + if self.num_locality_domains <= 1: + return None + locality_domain_id = self._next_locality_domain + self._next_locality_domain = (locality_domain_id + 1) % self.num_locality_domains + return locality_domain_id + + def get_locality_domain(self, request_id: int) -> int: + """Return the locality domain index assigned to an allocated request.""" + if self.num_locality_domains <= 1: + return 0 + kv_cache = self.kv_cache_map.get(request_id) + assert kv_cache is not None, ( + f"get_locality_domain: request {request_id} has no allocated KV cache " + f"(num_locality_domains={self.num_locality_domains}, " + f"kv_cache_map size={len(self.kv_cache_map)})" + ) + locality_domain_id = getattr(kv_cache, "locality_domain_id", None) + assert locality_domain_id is not None, ( + f"get_locality_domain: request {request_id} has no concrete locality_domain_id " + f"(num_locality_domains={self.num_locality_domains}, " + f"kv_cache_map size={len(self.kv_cache_map)})" + ) + return locality_domain_id + + def _request_locality_domain_id(self, req: LlmRequest) -> int | None: + """Return the request-owned locality domain id after validating the invariant.""" + locality_domain_id = getattr(req, "py_locality_domain_id", None) + if self.num_locality_domains <= 1: + assert locality_domain_id is None or locality_domain_id == 0, ( + "Non-localized requests must not use a nonzero locality_domain_id" + ) + return None + assert locality_domain_id is not None, ( + f"Request {req.py_request_id} must carry py_locality_domain_id " + "before localized KV-cache placement" + ) + assert 0 <= locality_domain_id < self.num_locality_domains, ( + f"locality_domain_id must be in [0, {self.num_locality_domains}), " + f"got {locality_domain_id}" + ) + return locality_domain_id + def get_num_available_tokens( self, *, token_num_upper_bound: int, batch_size: int = 1, max_num_draft_tokens: int = 0 ) -> int: @@ -4162,7 +4283,11 @@ def probe_context_reuse(self, req: LlmRequest) -> int | None: if not self.enable_block_reuse or req.py_request_id in self.kv_cache_map: return None return self.impl.probe_reuse( - ReuseScope(lora_id=req.lora_task_id, salt=self._derive_reuse_salt(req.cache_salt)), + ReuseScope( + lora_id=req.lora_task_id, + salt=self._derive_reuse_salt(req.cache_salt), + locality_domain_id=self._request_locality_domain_id(req), + ), self._context_reuse_tokens(req), ) @@ -4226,6 +4351,7 @@ def prepare_context_cache(self, req: LlmRequest, reuse_limit: int | None = None) # prompt_len - 1, because augmentation can change the length # for multimodal requests. num_lookup_tokens = len(tokens) if tokens is not None else None + locality_domain_id = self._request_locality_domain_id(req) kv_cache = self._create_kv_cache( req.py_request_id, req.lora_task_id, @@ -4237,9 +4363,14 @@ def prepare_context_cache(self, req: LlmRequest, reuse_limit: int | None = None) expected_prompt_length=( req.total_input_len_cp if self._has_cp_helix else req.prompt_len ), + locality_domain_id=locality_domain_id, ) if kv_cache is None: return None + if locality_domain_id is not None: + # Only localized placement has a domain to honor; the non-localized + # path leaves locality_domain_id None and has nothing to check. + assert getattr(kv_cache, "locality_domain_id", None) == locality_domain_id kv_cache.cuda_stream = self._cache_stream(req.py_request_id) if not self.enable_block_reuse: @@ -5089,6 +5220,7 @@ def _mirror_draft_kv_cache(self, req: LlmRequest): is_dummy=req.is_dummy, expected_prompt_length=req.prompt_len, kv_cache_retention_config=self._retention_config_for_request(req), + locality_domain_id=self._request_locality_domain_id(req), ) if kv_cache is None: return None @@ -5913,6 +6045,7 @@ def add_dummy_requests( encoder_output_lens: Optional[List[int]] = None, draft_kv_cache_manager: Optional["BaseResourceManager"] = None, capture_sampling_params: Optional[SamplingParams] = None, + locality_domain_ids: Optional[List[int | None]] = None, ): """Create dummy requests with allocated target and optional draft KV caches. @@ -5934,6 +6067,10 @@ def add_dummy_requests( on this manager and the supplied draft manager when it is no longer needed, including when subsequent warmup or graph capture fails. """ + if locality_domain_ids is not None: + assert len(locality_domain_ids) == len(request_ids), ( + "locality_domain_ids must have the same length as request_ids" + ) _kv_draft = ( kv_reserve_draft_tokens if kv_reserve_draft_tokens is not None else max_num_draft_tokens ) @@ -5988,6 +6125,11 @@ def release_resources( ) # Using 1 instead of 0 prevents NaN during warmup in e.g. Deepseek input_tokens = [1 for _ in range(token_num)] + request_locality_domain_id = None + if locality_domain_ids is not None: + request_locality_domain_id = locality_domain_ids[i] + if request_locality_domain_id is None and self.num_locality_domains > 1: + request_locality_domain_id = self.pick_locality_domain(req_id) req = LlmRequest( request_id=req_id, max_new_tokens=1, @@ -5995,6 +6137,7 @@ def release_resources( sampling_config=sampling_params._get_sampling_config(), is_streaming=False, encoder_input_tokens=encoder_input_tokens, + locality_domain_id=request_locality_domain_id, encoder_output_len=encoder_output_len, ) req.is_dummy_request = True @@ -6008,6 +6151,7 @@ def release_resources( input_tokens, is_dummy=req.is_dummy, expected_prompt_length=expected_prompt_length, + locality_domain_id=request_locality_domain_id, ) # Saturated IndexMapper (e.g. disagg gen trans in progress) # returns None; retry next iter. @@ -6030,16 +6174,21 @@ def release_resources( return None draft_kv_cache = None if draft_kv_cache_manager is not None: + draft_locality_domain_id = draft_kv_cache_manager._request_locality_domain_id(req) draft_kv_cache = draft_kv_cache_manager._create_kv_cache( req.py_request_id, req.lora_task_id, input_tokens, is_dummy=req.is_dummy, expected_prompt_length=expected_prompt_length, + locality_domain_id=draft_locality_domain_id, ) if draft_kv_cache is None: release_resources(req) return None + assert ( + getattr(draft_kv_cache, "locality_domain_id", None) == draft_locality_domain_id + ) success = draft_kv_cache.resume(draft_kv_cache_manager._stream.cuda_stream) if not success: release_resources(req, free_draft_resources=True) @@ -6448,16 +6597,65 @@ def calculate_scaling_factor_size_bytes( ) return get_size_in_bytes(cache_size // quant_vector_size, scaling_factor_dtype) + def _get_buffer_views_for_invalid_value_check( + self, layer_idx: int, kv_layout: str = "NHD" + ) -> list[torch.Tensor]: + """Return contiguous tensor views covering the mapped KV pages to scan.""" + if self.num_locality_domains <= 1: + return [self.get_buffers(layer_idx, kv_layout=kv_layout).flatten(0, 1)] + + layer_offset = self.layer_offsets[layer_idx] + element_per_container = 1 + dtype = self.dtype + if dtype == DataType.NVFP4: + element_per_container = 2 + dtype = torch.int8 + + views: list[torch.Tensor] = [] + for address, num_pages in self.impl.get_mapped_buffer_views( + LayerId(layer_offset), Role.KEY + ): + if kv_layout == "NHD": + shape = [ + num_pages // self.kv_factor, + self.kv_factor, + self.tokens_per_block, + self.num_kv_heads_per_layer[layer_offset], + self.head_dim_per_layer[layer_offset] // element_per_container, + ] + else: + shape = [ + num_pages // self.kv_factor, + self.kv_factor, + self.num_kv_heads_per_layer[layer_offset], + self.tokens_per_block, + self.head_dim_per_layer[layer_offset] // element_per_container, + ] + views.append( + convert_to_torch_tensor(TensorWrapper(int(address), dtype, shape)).flatten(0, 1) + ) + return views + def _iter_cache_buffers_for_invalid_check(self) -> Iterable[Tuple[int, torch.Tensor]]: pool_handled = set() for layer_id, layer_offset in self.layer_offsets.items(): pool_id = self.layer_to_pool_mapping_dict[layer_offset] if pool_id in pool_handled: continue - buffer = self.get_buffers(layer_id) - if buffer is None: - continue - yield layer_id, buffer + if self.num_locality_domains <= 1: + buffer = self.get_buffers(layer_id) + if buffer is None: + continue + yield layer_id, buffer + else: + # Guard pages are indexed against the whole layer buffer, which does + # not exist as one tensor once the pool is split, so yield no layer + # key: restoring the guard into a per-domain view would write the + # sentinel to the wrong page. + for buffer in self._get_buffer_views_for_invalid_value_check(layer_id): + if buffer is None: + continue + yield None, buffer pool_handled.add(pool_id) def check_invalid_values_in_kv_cache(self, fill_with_zero: bool = False) -> bool: @@ -7100,10 +7298,23 @@ def _create_kv_cache( enable_request_stats: bool = False, expected_prompt_length: int | None = None, kv_cache_retention_config: KvCacheRetentionConfig | None = None, + locality_domain_id: int | None = None, ): assert request_id not in self.kv_cache_map, ( f"KV cache for request {request_id} already exists" ) + if self.num_locality_domains > 1: + assert locality_domain_id is not None, ( + "Localized KV cache managers require a concrete locality_domain_id" + ) + assert 0 <= locality_domain_id < self.num_locality_domains, ( + f"locality_domain_id must be in [0, {self.num_locality_domains}), " + f"got {locality_domain_id}" + ) + else: + assert locality_domain_id is None, ( + "Non-localized KV cache managers require locality_domain_id=None" + ) if self.index_mapper.num_free_slots() == 0: logger.warning( "No free IndexMapper slots for request %s " @@ -7135,7 +7346,9 @@ def _create_kv_cache( expected_prompt_length, self.tokens_per_block, self._swa_endpoint_rewind ) kv_cache = self.impl.create_kv_cache( - ReuseScope(lora_id=lora_task_id, salt=reuse_salt), + ReuseScope( + lora_id=lora_task_id, salt=reuse_salt, locality_domain_id=locality_domain_id + ), input_tokens, id=request_id, enable_request_stats=enable_request_stats, @@ -7173,9 +7386,19 @@ def probe_prefix_match_length(self, input_tokens, lora_task_id=None, cache_salt= if len(input_tokens) == 0: return 0 reuse_salt = self._derive_reuse_salt(cache_salt) - return self.impl.probe_reuse( - ReuseScope(lora_id=lora_task_id, salt=reuse_salt), - input_tokens, + if self.num_locality_domains <= 1: + return self.impl.probe_reuse( + ReuseScope(lora_id=lora_task_id, salt=reuse_salt), + input_tokens, + ) + # The router has no request yet, so no domain: blocks are keyed per + # domain and the request may land in any of them. + return max( + self.impl.probe_reuse( + ReuseScope(lora_id=lora_task_id, salt=reuse_salt, locality_domain_id=domain), + input_tokens, + ) + for domain in range(self.num_locality_domains) ) def probe_first_new_block_key(self, req: LlmRequest) -> bytes | None: @@ -7218,7 +7441,11 @@ def probe_first_new_block_key(self, req: LlmRequest) -> bytes | None: # A 1-token prompt marshals to an empty sequence: the last token is # excluded, so there is nothing to look up and nothing to contribute. return None - scope = ReuseScope(lora_id=req.lora_task_id, salt=self._derive_reuse_salt(req.cache_salt)) + scope = ReuseScope( + lora_id=req.lora_task_id, + salt=self._derive_reuse_salt(req.cache_salt), + locality_domain_id=self._request_locality_domain_id(req), + ) return self.impl.probe_first_new_block_key(scope, tokens) def prefetch_for_context_tokens(self, requests: list) -> bool: @@ -7238,7 +7465,13 @@ def prefetch_for_context_tokens(self, requests: list) -> bool: # cache hits the same radix-tree blocks. reuse_salt = self._derive_reuse_salt(req.cache_salt) prefetched = self.impl.prefetch_reuse( - ReuseScope(lora_id=req.lora_task_id, salt=reuse_salt), tokens, CACHE_LEVEL1 + ReuseScope( + lora_id=req.lora_task_id, + salt=reuse_salt, + locality_domain_id=self._request_locality_domain_id(req), + ), + tokens, + CACHE_LEVEL1, ) if not prefetched: logger.warning("prefetch failed for request %s", req.py_request_id) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index df494fc9f27b..97638dc79ae7 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -3806,6 +3806,7 @@ def _create_kv_cache( enable_request_stats: bool = False, expected_prompt_length: Optional[int] = None, kv_cache_retention_config: Optional[KvCacheRetentionConfig] = None, + locality_domain_id: Optional[int] = None, ) -> Optional[_KVCache]: kv_cache = super()._create_kv_cache( request_id, @@ -3816,6 +3817,7 @@ def _create_kv_cache( enable_request_stats=enable_request_stats, expected_prompt_length=expected_prompt_length, kv_cache_retention_config=kv_cache_retention_config, + locality_domain_id=locality_domain_id, ) if (self.mapping.rank == 0 and kv_cache is not None and input_tokens is not None and not is_dummy @@ -4654,6 +4656,7 @@ def add_dummy_requests( encoder_output_lens: Optional[List[int]] = None, draft_kv_cache_manager: Optional[BaseResourceManager] = None, capture_sampling_params: Optional["SamplingParams"] = None, + locality_domain_ids: Optional[List[Optional[int]]] = None, ) -> List[LlmRequest]: requests = super().add_dummy_requests( request_ids=request_ids, @@ -4666,6 +4669,7 @@ def add_dummy_requests( encoder_output_lens=encoder_output_lens, draft_kv_cache_manager=draft_kv_cache_manager, capture_sampling_params=capture_sampling_params, + locality_domain_ids=locality_domain_ids, ) if requests: self._setup_state_indices( diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index e8375c682b8f..219e41d9509d 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -870,6 +870,7 @@ def __init__( logprobs_mode: LogprobMode = LogprobMode.RAW, logprobs_simple_format: bool = False, return_routed_experts: bool = False, + locality_domain_id: Optional[int] = None, **kwargs): self.py_logits_post_processors = kwargs.pop("py_logits_post_processors", None) @@ -950,6 +951,9 @@ def __init__( self.is_cuda_graph_dummy = False self.py_kv_transfer_start_time = None self.py_kv_transfer_timed_out = False + # Python-side locality domain ownership. ``None`` is used for non-localized + # execution; localized requests must pass a concrete value at creation. + self.py_locality_domain_id = locality_domain_id # Prevent recreation after a send session drops peer registration. self.py_kv_send_session_retired = False @@ -1628,7 +1632,8 @@ def executor_request_to_llm_request( child_req_ids: List[int], exclude_last_generation_logits: bool, input_token_ids: Optional[List] = None, - position_ids: Optional[List] = None) -> LlmRequest: + position_ids: Optional[List] = None, + locality_domain_id: Optional[int] = None) -> LlmRequest: sampling_config = executor_request.sampling_config input_tokens = input_token_ids if input_token_ids is not None else executor_request.input_token_ids @@ -1772,6 +1777,7 @@ def executor_request_to_llm_request( "py_logprobs_simple_format", False), return_routed_experts=getattr(executor_request, "py_return_routed_experts", False), + locality_domain_id=locality_domain_id, ) # Bad-words list for the TorchSampler path, kept in its native diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 0406c2d58e73..1ed2d18e3f05 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -6882,6 +6882,12 @@ def _fetch_new_requests( new_requests = new_requests_cur_rank + if self._is_kv_manager_v2: + for req_item in new_requests: + if req_item.is_normal_request: + req_item.locality_domain_id = ( + self.kv_cache_manager.pick_locality_domain(req_item.id)) + # 7. Merge requests if self._kv_load_tracker.enabled and new_requests: self._kv_load_tracker.idle_sampled = False diff --git a/tensorrt_llm/_torch/pyexecutor/request_utils.py b/tensorrt_llm/_torch/pyexecutor/request_utils.py index 4a61d96b0dd2..ae031ba048c8 100644 --- a/tensorrt_llm/_torch/pyexecutor/request_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/request_utils.py @@ -273,7 +273,11 @@ def merge_requests_to_llm_requests( req_with_children = [] for req_item in new_requests: req = executor_request_to_llm_request( - req_item.id, req_item.request, req_item.child_req_ids, exclude_last_generation_logits + req_item.id, + req_item.request, + req_item.child_req_ids, + exclude_last_generation_logits, + locality_domain_id=req_item.locality_domain_id, ) req_with_children.append(req) if req.child_requests: @@ -317,6 +321,7 @@ def merge_helix_requests( exclude_last_generation_logits=exclude_last_generation_logits, input_token_ids=input_ids_this_rank, position_ids=position_ids_this_rank, + locality_domain_id=req_item.locality_domain_id, ) req.total_input_len_cp = input_len req.seqlen_this_rank_cp = len(input_ids_this_rank) diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index bf4c0d879fd5..dd5d9c0c9ccb 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -544,7 +544,12 @@ def preempt_for_pages(req: LlmRequest) -> bool: protected.update(r.py_request_id for r in scheduled_ctx) protected.add(req.py_request_id) return self._try_preempt_for_pages( - requests_list, protected, inflight_request_ids, recompute_paused, preempted_ids + requests_list, + protected, + inflight_request_ids, + recompute_paused, + preempted_ids, + locality_domain_id=getattr(req, "py_locality_domain_id", None), ) # --- Phase 2: schedule deferred context / encoder requests --- @@ -1614,6 +1619,7 @@ def _try_preempt_for_pages( inflight_request_ids: set[int], recompute_paused: RequestList, preempted_ids: set[int], + locality_domain_id: int | None = None, ) -> bool: """Release one started request's KV cache so another can allocate. @@ -1643,6 +1649,15 @@ def _try_preempt_for_pages( victim = requests_list[i] if victim.py_request_id in protected_ids: continue + # Pages are allocated per locality domain, so a victim in another + # domain frees nothing for this request. + victim_domain = getattr(victim, "py_locality_domain_id", None) + if ( + isinstance(locality_domain_id, int) + and isinstance(victim_domain, int) + and victim_domain != locality_domain_id + ): + continue if not self._is_recompute_pause_candidate(victim, inflight_request_ids): continue if not self.kv_cache_manager.is_request_active(victim.py_request_id): diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py index 255bc46ff80c..62f7237d75e0 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py @@ -74,7 +74,9 @@ def _make_creator( c._speculative_config = None c._mapping = Mock() c._model_engine = Mock() - c._llm_args = SimpleNamespace(kv_cache_compression_config=None, sleep_config=None) + c._llm_args = SimpleNamespace( + kv_cache_compression_config=None, sleep_config=None, enable_locality_domains=False + ) c._disable_overlap_scheduler = False c._kv_cache_manager_cls = Mock() @@ -369,6 +371,9 @@ def get_cache_size_per_token( creator._kv_cache_manager_cls = RecordingKVCacheManager creator._is_disagg = False creator._cache_transceiver_config = None + creator._llm_args = SimpleNamespace( + kv_cache_compression_config=None, sleep_config=None, enable_locality_domains=False + ) creator._should_create_separate_draft_kv_cache = Mock(return_value=True) creator._get_effective_draft_config = Mock(return_value=draft_model_config) creator._get_num_draft_layers = Mock(return_value=1) @@ -1078,7 +1083,9 @@ def get_cache_size_per_token(model_config, *args, **kwargs): c._max_num_tokens = 8192 c._max_beam_width = 1 c._cache_transceiver_config = None - c._llm_args = SimpleNamespace(kv_cache_compression_config=None, sleep_config=None) + c._llm_args = SimpleNamespace( + kv_cache_compression_config=None, sleep_config=None, enable_locality_domains=False + ) c._mapping = Mock(enable_attention_dp=True, tp_size=1) c._mapping.has_cp_helix.return_value = False c._mapping.pp_layers.return_value = list(range(self.DRAFT_LAYERS)) diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py index 116d43a6849a..06e530ccbe53 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py @@ -1678,6 +1678,7 @@ def test_separate_one_model_draft_normalizes_target_pool_ratio() -> None: creator._mapping = Mock() creator._speculative_config = Mock() creator._disable_overlap_scheduler = False + creator._llm_args = SimpleNamespace(enable_locality_domains=False) effective_draft_config = Mock() effective_draft_config.pretrained_config.torch_dtype = "bfloat16" diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py index 0abaa5a724ab..28dca0cedd3c 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py @@ -2136,6 +2136,24 @@ def test_freed_pages_are_reserved_for_the_request_that_preempted(self): assert ids(out.recompute_paused_requests) == [99] assert ids(out.context_requests) == [] + def test_preemption_victim_must_share_the_locality_domain(self): + """Pages are per domain: a victim elsewhere frees nothing for the request.""" + mgr = make_kv_cache_manager(has_cache_tier_below_gpu=False) + sched = make_scheduler(mgr, max_num_tokens=1000) + same = make_gen_request(1) + same.py_locality_domain_id = 0 + # Newest, so it would be taken first without the domain filter. + other = make_gen_request(2) + other.py_locality_domain_id = 1 + recompute_paused = [] + + assert sched._try_preempt_for_pages( + [same, other], set(), set(), recompute_paused, set(), locality_domain_id=0 + ) + + assert recompute_paused == [same] + mgr.preempt_request.assert_called_once_with(same) + def test_a_failed_preemption_still_lets_the_pass_continue(self): """With nothing to give up there is no capacity to reserve.""" mgr = make_kv_cache_manager( diff --git a/tests/unittest/_torch/executor/test_locality_domain_request.py b/tests/unittest/_torch/executor/test_locality_domain_request.py new file mode 100644 index 000000000000..9e7010646c1a --- /dev/null +++ b/tests/unittest/_torch/executor/test_locality_domain_request.py @@ -0,0 +1,84 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Locality domain plumbing that needs no GPU: request propagation and fork-join gating.""" + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( + KVCacheManagerV2, + resolve_fork_join_attn, +) +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest +from tensorrt_llm.sampling_params import SamplingParams + +pytestmark = pytest.mark.cpu_only + + +def _request(request_id: int, locality_domain_id): + return LlmRequest( + request_id=request_id, + max_new_tokens=4, + input_tokens=[1, 2, 3], + sampling_config=SamplingParams(n=2, best_of=2, top_k=4)._get_sampling_config(), + is_streaming=False, + locality_domain_id=locality_domain_id, + ) + + +@pytest.mark.parametrize("domain", [None, 0, 1]) +def test_child_request_inherits_locality_domain(domain): + parent = _request(7, domain) + assert parent.py_locality_domain_id == domain + parent.create_child_request(8) + (child,) = parent.child_requests + assert child.py_locality_domain_id == domain + + +def test_fork_join_requires_two_domains_and_the_runtime(): + assert resolve_fork_join_attn(2, True, runtime_enabled=True, mock_localization=False) + assert not resolve_fork_join_attn(1, True, runtime_enabled=True, mock_localization=False) + assert not resolve_fork_join_attn(2, False, runtime_enabled=True, mock_localization=False) + # Mock localization keeps the split KV cache on one stream. + assert not resolve_fork_join_attn(2, True, runtime_enabled=False, mock_localization=True) + + +def test_fork_join_rejects_bad_configurations(): + with pytest.raises(ValueError, match="exactly 2"): + resolve_fork_join_attn(3, True, runtime_enabled=True, mock_localization=False) + with pytest.raises(RuntimeError, match="disabled on this device"): + resolve_fork_join_attn(2, True, runtime_enabled=False, mock_localization=False) + + +def test_locality_domains_reject_sparse_attention_at_manager_creation(): + from tensorrt_llm._torch.pyexecutor._util import KvCacheCreator + + creator = KvCacheCreator.__new__(KvCacheCreator) + creator._llm_args = SimpleNamespace(enable_locality_domains=True) + creator._sparse_attention_config = SimpleNamespace(algorithm="dsa") + with pytest.raises(ValueError, match="sparse attention"): + creator._create_kv_cache_manager(Mock()) + + +def test_placement_alternates_domains_when_ids_advance_in_even_steps(): + # n=2 sampling gives each request a parent and a child id, so new requests + # arrive with ids 0, 2, 4, ...; parity placement would put them all on domain 0. + manager = SimpleNamespace(num_locality_domains=2, _next_locality_domain=0) + picks = [ + KVCacheManagerV2.pick_locality_domain(manager, request_id) for request_id in (0, 2, 4, 6) + ] + assert picks == [0, 1, 0, 1] diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 285b8ad28ecf..27d7ef9996e1 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -2929,6 +2929,7 @@ def _run_non_adp_fetch(active_requests, *, exclude_retiring_requests, max_num_ac _should_exclude_last_generation_logits=lambda: False, _kv_load_tracker=KvLoadTracker(None, rank=0, enabled=False), _kv_cache_capacity={}, + _is_kv_manager_v2=False, ) seen = {} stub._fetch_and_enqueue_requests = lambda waiting_queue, live: seen.update(live=live) @@ -3014,6 +3015,7 @@ def test_non_adp_fetch_pops_a_full_batch_beside_retiring_requests( _should_exclude_last_generation_logits=lambda: False, _kv_load_tracker=KvLoadTracker(None, rank=0, enabled=False), _kv_cache_capacity={}, + _is_kv_manager_v2=False, ) for name in ("_get_request_admission_capacity", "_pop_from_waiting_queue"): setattr(stub, name, types.MethodType(getattr(PyExecutor, name), stub)) diff --git a/tests/unittest/_torch/executor/test_seq_slot_sizing.py b/tests/unittest/_torch/executor/test_seq_slot_sizing.py index deab2eae9e92..1707b4b44588 100644 --- a/tests/unittest/_torch/executor/test_seq_slot_sizing.py +++ b/tests/unittest/_torch/executor/test_seq_slot_sizing.py @@ -769,6 +769,7 @@ def _make_kv_cache_creator(disable_overlap_scheduler: bool, is_v2: bool = True) c._disable_overlap_scheduler = disable_overlap_scheduler c._llm_args = SimpleNamespace( kv_cache_config=SimpleNamespace(kv_events_config=None), + enable_locality_domains=False, ) # Short-circuit the post-construction max_seq_len fixup. c._skip_est = True diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index fbbafa3e7d68..71355974ab77 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -406,6 +406,23 @@ def test_from_yaml_migrates_legacy_mamba_interval(self, tmp_path): assert llm_args.kv_cache_config.mamba_state_config.periodic_snapshot_interval == 64 + def test_from_yaml_enable_locality_domains(self, tmp_path): + assert TorchLlmArgs( + model=llama_model_path).enable_locality_domains is False + + yaml_path = tmp_path / "locality_domains.yaml" + yaml_path.write_text( + yaml.safe_dump({ + "model": str(llama_model_path), + "enable_locality_domains": True, + }), + encoding="utf-8", + ) + + llm_args = TorchLlmArgs.from_yaml(yaml_path) + + assert llm_args.enable_locality_domains is True + def test_from_yaml_empty_file_reports_missing_model(self, tmp_path): yaml_path = tmp_path / "empty.yaml" yaml_path.write_text("", encoding="utf-8") From 09fd4a5259db414eb269c312348507fd5fab20c6 Mon Sep 17 00:00:00 2001 From: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> Date: Fri, 23 Oct 2026 18:38:48 +0000 Subject: [PATCH 3/5] [None][feat] Run attention once per locality domain on that domain's stream With locality domains, attention forks one call per domain onto the domain's Green Context stream, with metadata sliced to that domain's requests, and joins afterwards. The decoder runner groups generation rows by domain (keeping the overlap scheduler's host-token rows first), splits by ownership when the rows are domain-sorted and by position otherwise, and builds per-domain warmup and CUDA graph dummies. Per-domain MLA decode state uses one counter allocation and reuses domain 0's sequence offsets. Speculative decoding slices its per-domain offsets. Partitioned attention rejects FlashMLA, beam search, cross attention, helix, custom masks, sparse indices and MRoPE outside the mixed input type. Signed-off-by: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> --- cpp/tensorrt_llm/nanobind/thop/bindings.cpp | 12 +- cpp/tensorrt_llm/thop/attentionOp.cpp | 16 +- cpp/tensorrt_llm/thop/attentionOp.h | 2 +- cpp/tensorrt_llm/thop/cublasScaledMM.cpp | 41 +- .../attention/backends/fmha/fallback.py | 1 + .../_torch/attention/backends/interface.py | 3 + .../_torch/attention/backends/trtllm.py | 1169 ++++++++++++++++- .../engine/runners/decoder/dummy_requests.py | 153 ++- .../engine/runners/decoder/graphs.py | 4 + .../engine/runners/decoder/inputs.py | 80 ++ .../engine/runners/decoder/locality_domain.py | 99 ++ tensorrt_llm/_torch/speculative/eagle3.py | 6 +- tensorrt_llm/_torch/speculative/interface.py | 61 +- tensorrt_llm/_torch/speculative/mtp.py | 5 +- .../kv_cache/test_stream_localization_e2e.py | 99 ++ .../_torch/attention/test_attention_mla.py | 5 + .../test_locality_domain_attention_gpu.py | 306 +++++ .../test_locality_domain_partitions.py | 435 ++++++ .../executor/test_context_only_warmup.py | 1 + .../executor/test_locality_domain_dispatch.py | 114 ++ .../executor/test_pytorch_model_engine.py | 55 +- 21 files changed, 2566 insertions(+), 101 deletions(-) create mode 100644 tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/locality_domain.py create mode 100644 tests/integration/defs/kv_cache/test_stream_localization_e2e.py create mode 100644 tests/unittest/_torch/attention/test_locality_domain_attention_gpu.py create mode 100644 tests/unittest/_torch/attention/test_locality_domain_partitions.py create mode 100644 tests/unittest/_torch/executor/test_locality_domain_dispatch.py diff --git a/cpp/tensorrt_llm/nanobind/thop/bindings.cpp b/cpp/tensorrt_llm/nanobind/thop/bindings.cpp index a41158dd4661..9b6ba1f45150 100644 --- a/cpp/tensorrt_llm/nanobind/thop/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/thop/bindings.cpp @@ -260,12 +260,12 @@ void initBindings(nb::module_& m) nb::arg("aux_kv_cache_pool_ptr") = std::nullopt, nb::arg("is_cross") = false, nb::arg("cross_kv") = std::nullopt, nb::arg("relative_attention_bias") = std::nullopt, nb::arg("relative_attention_max_distance") = 0, nb::arg("spec_decoding_target_max_draft_tokens") = std::nullopt, - nb::arg("quant_scale_qkv") = std::nullopt, nb::arg("dsv4_inv_rope_cos_sin_cache") = std::nullopt, - nb::arg("enable_dsv4_epilogue_fusion") = false, nb::arg("force_prepare_spec_dec_tree_mask") = false, - nb::arg("max_num_sequences") = std::nullopt, nb::arg("kv_norm_weight") = std::nullopt, - nb::arg("kv_norm_eps") = 1e-6, nb::arg("skip_correction_threshold") = 0.0, - nb::arg("uses_spcompress") = std::nullopt, "Multi-head attention operation", - nb::call_guard()); + nb::arg("locality_domain_id") = -1, nb::arg("quant_scale_qkv") = std::nullopt, + nb::arg("dsv4_inv_rope_cos_sin_cache") = std::nullopt, nb::arg("enable_dsv4_epilogue_fusion") = false, + nb::arg("force_prepare_spec_dec_tree_mask") = false, nb::arg("max_num_sequences") = std::nullopt, + nb::arg("kv_norm_weight") = std::nullopt, nb::arg("kv_norm_eps") = 1e-6, + nb::arg("skip_correction_threshold") = 0.0, nb::arg("uses_spcompress") = std::nullopt, + "Multi-head attention operation", nb::call_guard()); m.def( "get_helix_workspace_size_per_rank", diff --git a/cpp/tensorrt_llm/thop/attentionOp.cpp b/cpp/tensorrt_llm/thop/attentionOp.cpp index 25c37e08fb91..852bc404a8d2 100644 --- a/cpp/tensorrt_llm/thop/attentionOp.cpp +++ b/cpp/tensorrt_llm/thop/attentionOp.cpp @@ -1135,9 +1135,10 @@ using torch_ext::trtllm::attention::Runner; using torch_ext::trtllm::attention::AttentionInputType; static std::shared_ptr get_attention_op( - RunnerPtr const& runner, std::shared_ptr& op, int64_t local_layer_idx) + RunnerPtr const& runner, std::shared_ptr& op, int64_t local_layer_idx, int64_t locality_domain_id) { - auto cache_key = std::make_tuple(op->data(), runner->data()); + // Each locality domain runs on its own stream and must not share an op. + auto cache_key = std::make_tuple(op->data(), runner->data(), locality_domain_id); using CacheKey = decltype(cache_key); static std::unordered_map, OpCustomHash> op_cache; static std::shared_mutex op_cache_mutex; @@ -1237,10 +1238,11 @@ void attention(torch::Tensor q, std::optional k, std::optional aux_kv_cache_pool_ptr, bool const is_cross, std::optional cross_kv, std::optional relative_attention_bias, int64_t relative_attention_max_distance, std::optional spec_decoding_target_max_draft_tokens, - std::optional quant_scale_qkv, std::optional dsv4_inv_rope_cos_sin_cache, - bool enable_dsv4_epilogue_fusion, bool const force_prepare_spec_dec_tree_mask, - std::optional const max_num_sequences, std::optional kv_norm_weight, double kv_norm_eps, - double skip_correction_threshold, std::optional uses_spcompress) + int64_t locality_domain_id, std::optional quant_scale_qkv, + std::optional dsv4_inv_rope_cos_sin_cache, bool enable_dsv4_epilogue_fusion, + bool const force_prepare_spec_dec_tree_mask, std::optional const max_num_sequences, + std::optional kv_norm_weight, double kv_norm_eps, double skip_correction_threshold, + std::optional uses_spcompress) { TLLM_LOG_TRACE("Attention op starts at layer %d", local_layer_idx); // Use these tensors to infer if the attention is using KV cache @@ -1485,7 +1487,7 @@ void attention(torch::Tensor q, std::optional k, std::optionalmFP8ContextFMHA || op->mFP8ContextMLA, "uses_spcompress requires FP8 context FMHA or FP8 context MLA."); } - op = get_attention_op(runner, op, local_layer_idx); + op = get_attention_op(runner, op, local_layer_idx, locality_domain_id); int32_t const num_seqs = host_context_lengths.size(0); RequestType const* request_types = static_cast(host_request_types.data_ptr()); diff --git a/cpp/tensorrt_llm/thop/attentionOp.h b/cpp/tensorrt_llm/thop/attentionOp.h index 1c12001c226a..9e0a7fd134f0 100644 --- a/cpp/tensorrt_llm/thop/attentionOp.h +++ b/cpp/tensorrt_llm/thop/attentionOp.h @@ -92,7 +92,7 @@ void attention(torch::Tensor q, std::optional k, std::optional aux_kv_cache_pool_ptr = std::nullopt, bool const is_cross = false, std::optional cross_kv = std::nullopt, std::optional relative_attention_bias = std::nullopt, int64_t relative_attention_max_distance = 0, - std::optional spec_decoding_target_max_draft_tokens = std::nullopt, + std::optional spec_decoding_target_max_draft_tokens = std::nullopt, int64_t locality_domain_id = -1, std::optional quant_scale_qkv = std::nullopt, std::optional dsv4_inv_rope_cos_sin_cache = std::nullopt, bool enable_dsv4_epilogue_fusion = false, bool const force_prepare_spec_dec_tree_mask = false, std::optional const max_num_sequences = std::nullopt, diff --git a/cpp/tensorrt_llm/thop/cublasScaledMM.cpp b/cpp/tensorrt_llm/thop/cublasScaledMM.cpp index 62f51f7b06f9..83d3255843e8 100644 --- a/cpp/tensorrt_llm/thop/cublasScaledMM.cpp +++ b/cpp/tensorrt_llm/thop/cublasScaledMM.cpp @@ -22,7 +22,9 @@ #include "tensorrt_llm/thop/outputTensor.h" #include "tensorrt_llm/thop/thUtils.h" #include "userbuffersTensor.h" +#include #include +#include #include #include @@ -41,6 +43,36 @@ using tensorrt_llm::common::CublasMMWrapper; using cublas_lut::AlgoListType; +CUcontext getCurrentCudaContext() +{ + CUcontext context{}; + CUresult result = cuCtxGetCurrent(&context); + if (result == CUDA_ERROR_NOT_INITIALIZED || context == nullptr) + { + TLLM_CUDA_CHECK(cudaFree(nullptr)); + result = cuCtxGetCurrent(&context); + } + TLLM_CHECK_WITH_INFO(result == CUDA_SUCCESS, "cuCtxGetCurrent failed with error %d", result); + TLLM_CHECK_WITH_INFO(context != nullptr, "Current CUDA context is null"); + return context; +} + +// Each locality domain runs under its own green CUDA context, so the wrapper +// cannot be a single thread-local instance. +std::shared_ptr getCublasWrapper() +{ + thread_local std::unordered_map> cublasWrappers; + auto& cublasWrapper = cublasWrappers[getCurrentCudaContext()]; + if (!cublasWrapper) + { + auto cublasHandle = getCublasHandle(); + auto cublasLtHandle = getCublasLtHandle(); + cublasWrapper = std::make_shared(cublasHandle, cublasLtHandle, nullptr, nullptr); + } + + return cublasWrapper; +} + void set_algo_attr(cublasLtMatmulAlgo_t& algo, std::array const& attr_list) { auto const& [algoId, tileID, stagesID, numsK, reduction, swizzle, customOption_, cga_] = attr_list; @@ -198,13 +230,8 @@ void cublas_gemm_caller(torch::Tensor& out, torch::Tensor const& a, torch::Tenso int32_t n = b.sizes()[1]; int32_t k = a.sizes()[1]; - thread_local std::shared_ptr cublasWrapper; - if (cublasWrapper == nullptr) - { - auto cublasHandle = getCublasHandle(); - auto cublasLtHandle = getCublasLtHandle(); - cublasWrapper = std::make_shared(cublasHandle, cublasLtHandle, nullptr, nullptr); - } + at::cuda::CUDAGuard deviceGuard(a.device()); + auto cublasWrapper = getCublasWrapper(); cudaDataType_t aType = convert_torch_dtype(a.scalar_type()); cudaDataType_t bType = convert_torch_dtype(b.scalar_type()); diff --git a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py index 0831ace484f8..1c0a9c28ae5a 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py @@ -417,6 +417,7 @@ def forward( v_head_dim=attn.v_head_dim, rope_append=attn.rope_append, attention_chunk_size=attn.attention_chunk_size, + locality_domain_id=forward_args.locality_domain_id, skip_softmax_stat=attn.skip_softmax_stat, skip_correction_threshold=attn.skip_correction_threshold, uses_spcompress=attn.uses_spcompress, diff --git a/tensorrt_llm/_torch/attention/backends/interface.py b/tensorrt_llm/_torch/attention/backends/interface.py index 4b411b09edc8..e903e92ccc4e 100644 --- a/tensorrt_llm/_torch/attention/backends/interface.py +++ b/tensorrt_llm/_torch/attention/backends/interface.py @@ -1032,6 +1032,9 @@ class AttentionForwardArgs: cu_q_seqlens: Optional[torch.Tensor] = None cu_kv_seqlens: Optional[torch.Tensor] = None fmha_scheduler_counter: Optional[torch.Tensor] = None + + # Locality domain this call runs against; -1 means the whole batch. + locality_domain_id: int = -1 # Testing only: skip the RoPE step of MLA generation (the standalone harness # feeds a pre-RoPE'd fused_q). The TRTLLM backend then appends the new latent # and inits the trtllm-gen scheduler buffers itself. diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index d8722ce74b67..f816db59951d 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -18,7 +18,8 @@ import math import os import weakref -from dataclasses import dataclass, field +from contextlib import contextmanager +from dataclasses import dataclass, field, replace from typing import (TYPE_CHECKING, Any, ClassVar, Dict, FrozenSet, Iterator, List, Optional, Tuple) @@ -33,6 +34,8 @@ from ...speculative.spec_tree_manager import SpecTreeManager from .sparse.dsa.metadata import TokenMajorGenView +from itertools import accumulate + from tensorrt_llm._utils import get_sm_version, maybe_pin_memory, prefer_pinned from tensorrt_llm.bindings.internal import thop from tensorrt_llm.functional import AttentionMaskType @@ -40,6 +43,7 @@ from tensorrt_llm.math_utils import ceil_div from tensorrt_llm.models.modeling_utils import QuantConfig +from ...locality_domain.runtime import LocalityDomainRuntime from ...pyexecutor.config_utils import is_mla from ...utils import (compute_swizzled_sf_shape, get_global_attrs, get_model_extra_attrs, helix_local_len_tensor) @@ -60,6 +64,12 @@ _SKIP_CORRECTION_SUPPORTED_SMS = frozenset((100, 103)) +# Set to 0 to run the locality domains one at a time instead of concurrently. +# Each domain still uses its own stream, workspace and metadata, so only the +# overlap is lost -- useful for telling a partitioning bug apart from a race. +_LOCALITY_DOMAIN_TWO_STREAM = os.environ.get( + "TRTLLM_LOCALITY_DOMAIN_TWO_STREAM_ATTN", "1") == "1" + def _resolve_skip_correction_threshold(threshold: float, sm_version: int, @@ -315,6 +325,34 @@ def max_context_length(self) -> int: """ return min(self.max_seq_len, self.max_num_tokens) + # Locality domain two-kernel attention fields (backend-specific, not in interface.py) + locality_domain_enabled: bool = False + locality_domain_runtime: Optional[LocalityDomainRuntime] = None + locality_domain_num_ctx_tokens: Optional[List[int]] = None + locality_domain_num_gen_tokens: Optional[List[int]] = None + # Cumulative boundary arrays of length N+1 (N = num_locality_domains). + # locality_domain_ctx_req_splits[i]..[i+1] is the range of context requests + # belonging to locality domain i. + locality_domain_ctx_req_splits: Optional[List[int]] = None + locality_domain_gen_req_splits: Optional[List[int]] = None + locality_domain_ctx_token_splits: Optional[List[int]] = None + locality_domain_gen_token_splits: Optional[List[int]] = None + locality_domain_kv_cache_block_offsets: Optional[List[torch.Tensor]] = None + draft_locality_domain_kv_cache_block_offsets: Optional[List[ + torch.Tensor]] = None + # Per-locality-domain MLA generation state computed by mla_rope_generation. + locality_domain_mla_gen_state: Optional[List[Optional[dict]]] = None + # Pre-allocated per-locality-domain workspaces. + locality_domain_workspaces: Optional[List[torch.Tensor]] = None + locality_domain_cuda_graph_workspaces: Optional[List[torch.Tensor]] = None + # Per-domain trtllm-gen tree mask buffers; entry 0 is None (domain 0 keeps + # spec_decoding_bl_tree_mask). + locality_domain_spec_decoding_bl_tree_masks: Optional[List[Optional[ + torch.Tensor]]] = None + # Layout whose domain slices last passed _validate_locality_domain_coverage; + # the check reruns only when the layout changes. + locality_domain_validated_layout: Optional[tuple] = None + @property def max_seq_len(self) -> int: """ @@ -481,6 +519,8 @@ def _post_init_with_buffers(self, buffers) -> None: dtype=torch.int8, ) + self._ensure_locality_domain_workspaces() + if self.kv_cache_manager is not None: num_attention_op_pools = getattr(self.kv_cache_manager, "num_attention_op_pools", @@ -694,6 +734,27 @@ def on_update_kv_lens(self): if state is not None: state.on_update_kv_lens(self) + def _ensure_locality_domain_workspaces(self) -> None: + if not self.locality_domain_enabled: + return + + num_locality_domains = (self.kv_cache_manager.num_locality_domains + if self.kv_cache_manager is not None else 2) + if (not self.is_cuda_graph and + (self.locality_domain_workspaces is None + or len(self.locality_domain_workspaces) != num_locality_domains)): + self.locality_domain_workspaces = [ + torch.empty(0, device='cuda', dtype=torch.int8) + for _ in range(num_locality_domains) + ] + if (self.locality_domain_cuda_graph_workspaces is None + or len(self.locality_domain_cuda_graph_workspaces) + != num_locality_domains): + self.locality_domain_cuda_graph_workspaces = [ + torch.empty(0, device='cuda', dtype=torch.int8) + for _ in range(num_locality_domains) + ] + def update_for_spec_dec(self) -> None: if self.enable_flash_mla: self._flash_mla_metadata_valid = False @@ -709,6 +770,92 @@ def _invalidate_mla_scheduler_buffers(self) -> None: self._mla_ctx_cu_seqlens_valid = False self._cute_dsl_mla_split_plan_key = None + if (not self.locality_domain_enabled or self.kv_cache_manager is None + or self.request_ids is None): + return + + num_locality_domains = self.kv_cache_manager.num_locality_domains + sequence_lengths = self.seq_lens[:self.num_seqs].tolist() + + per_locality_domain_ctx_reqs = [0] * num_locality_domains + per_locality_domain_ctx_tokens = [0] * num_locality_domains + per_locality_domain_gen_reqs = [0] * num_locality_domains + per_locality_domain_gen_tokens = [0] * num_locality_domains + + def locality_domain_from_splits(splits: Optional[List[int]], + local_req_idx: int) -> Optional[int]: + if splits is None: + return None + for locality_domain_idx in range(len(splits) - 1): + if splits[locality_domain_idx] <= local_req_idx < splits[ + locality_domain_idx + 1]: + return locality_domain_idx + return None + + old_ctx_req_splits = self.locality_domain_ctx_req_splits + old_gen_req_splits = self.locality_domain_gen_req_splits + for req_idx, request_id in enumerate(self.request_ids[:self.num_seqs]): + if req_idx < self.num_contexts: + locality_domain_id = locality_domain_from_splits( + old_ctx_req_splits, req_idx) + if locality_domain_id is None: + locality_domain_id = self.kv_cache_manager.get_locality_domain( + request_id) + per_locality_domain_ctx_reqs[locality_domain_id] += 1 + per_locality_domain_ctx_tokens[ + locality_domain_id] += sequence_lengths[req_idx] + else: + gen_req_idx = req_idx - self.num_contexts + locality_domain_id = locality_domain_from_splits( + old_gen_req_splits, gen_req_idx) + if locality_domain_id is None: + locality_domain_id = self.kv_cache_manager.get_locality_domain( + request_id) + per_locality_domain_gen_reqs[locality_domain_id] += 1 + per_locality_domain_gen_tokens[ + locality_domain_id] += sequence_lengths[req_idx] + self.locality_domain_ctx_req_splits = [0] + list( + accumulate(per_locality_domain_ctx_reqs)) + self.locality_domain_gen_req_splits = [0] + list( + accumulate(per_locality_domain_gen_reqs)) + self.locality_domain_ctx_token_splits = [0] + list( + accumulate(per_locality_domain_ctx_tokens)) + self.locality_domain_gen_token_splits = [0] + list( + accumulate(per_locality_domain_gen_tokens)) + self.locality_domain_num_ctx_tokens = per_locality_domain_ctx_tokens + self.locality_domain_num_gen_tokens = per_locality_domain_gen_tokens + + if self.host_request_types is not None: + self.host_request_types[:self.num_contexts].fill_(0) + self.host_request_types[self.num_contexts:self.num_seqs].fill_(1) + + if self.kv_cache_block_offsets is not None: + self.locality_domain_kv_cache_block_offsets = ( + self.build_locality_domain_block_offsets( + self.kv_cache_block_offsets, + '_locality_domain_block_offsets_buf')) + + if self.draft_kv_cache_block_offsets is not None: + self.draft_locality_domain_kv_cache_block_offsets = ( + self.build_locality_domain_block_offsets( + self.draft_kv_cache_block_offsets, + '_draft_locality_domain_block_offsets_buf')) + + if not hasattr(self, '_locality_domain_host_total_kv_lens_buf'): + self._locality_domain_host_total_kv_lens_buf = [ + torch.zeros(2, dtype=torch.int, device='cpu') + for _ in range(num_locality_domains) + ] + + if self.is_cuda_graph: + for locality_domain_idx in range(num_locality_domains): + buf = self._locality_domain_host_total_kv_lens_buf[ + locality_domain_idx] + buf[0] = self.max_seq_len * per_locality_domain_ctx_reqs[ + locality_domain_idx] + buf[1] = self.max_seq_len * per_locality_domain_gen_reqs[ + locality_domain_idx] + def restore_from_spec_dec(self) -> None: super().restore_from_spec_dec() if self.fp4_mla_state is not None: @@ -819,6 +966,51 @@ def recompute_helix_spec_buffers(self, num_gen_tokens: int, cp_size=self.mapping.cp_size, cp_rank=self.mapping.cp_rank) + def build_locality_domain_block_offsets( + self, + source_block_offsets: Optional[torch.Tensor], + buffer_attr: str, + ) -> Optional[List[torch.Tensor]]: + if (not self.locality_domain_enabled or source_block_offsets is None + or self.kv_cache_manager is None): + return None + + num_locality_domains = self.kv_cache_manager.num_locality_domains + buffers = getattr(self, buffer_attr, None) + if buffers is None: + buffers = [ + torch.zeros_like(source_block_offsets) + for _ in range(num_locality_domains) + ] + setattr(self, buffer_attr, buffers) + + packed_block_offsets = [] + num_ctx = self.num_contexts + for locality_domain_idx in range(num_locality_domains): + ctx_start = self.locality_domain_ctx_req_splits[locality_domain_idx] + ctx_end = self.locality_domain_ctx_req_splits[locality_domain_idx + + 1] + gen_start = num_ctx + self.locality_domain_gen_req_splits[ + locality_domain_idx] + gen_end = num_ctx + self.locality_domain_gen_req_splits[ + locality_domain_idx + 1] + num_ctx_reqs = ctx_end - ctx_start + num_gen_reqs = gen_end - gen_start + num_seqs = num_ctx_reqs + num_gen_reqs + block_offsets = buffers[locality_domain_idx] + block_offsets.zero_() + if num_seqs > 0: + if num_ctx_reqs > 0: + block_offsets[:, :num_ctx_reqs].copy_( + source_block_offsets[:, ctx_start:ctx_end], + non_blocking=True) + if num_gen_reqs > 0: + block_offsets[:, num_ctx_reqs:num_seqs].copy_( + source_block_offsets[:, gen_start:gen_end], + non_blocking=True) + packed_block_offsets.append(block_offsets) + return packed_block_offsets + def _bind_runtime_views( self, *, @@ -1124,6 +1316,64 @@ def prepare(self) -> None: self.num_seqs, max_blocks=max_blocks) + # Build per-locality-domain block offsets for fork-join attention. + if self.locality_domain_enabled: + self._ensure_locality_domain_workspaces() + num_locality_domains = self.kv_cache_manager.num_locality_domains + + if not hasattr(self, '_locality_domain_host_total_kv_lens_buf'): + self._locality_domain_host_total_kv_lens_buf = [ + torch.zeros(2, dtype=torch.int, device='cpu') + for _ in range(num_locality_domains) + ] + + self.locality_domain_kv_cache_block_offsets = ( + self.build_locality_domain_block_offsets( + self.kv_cache_block_offsets, + '_locality_domain_block_offsets_buf')) + + if (self.draft_kv_cache_manager is not None and getattr( + self.draft_kv_cache_manager, 'fork_join_attn', False)): + self.draft_locality_domain_kv_cache_block_offsets = ( + self.build_locality_domain_block_offsets( + self.draft_kv_cache_block_offsets, + '_draft_locality_domain_block_offsets_buf')) + + num_ctx = self.num_contexts + for locality_domain_idx in range(num_locality_domains): + ctx_start = self.locality_domain_ctx_req_splits[ + locality_domain_idx] + ctx_end = self.locality_domain_ctx_req_splits[ + locality_domain_idx + 1] + gen_start = num_ctx + self.locality_domain_gen_req_splits[ + locality_domain_idx] + gen_end = num_ctx + self.locality_domain_gen_req_splits[ + locality_domain_idx + 1] + num_ctx_reqs = ctx_end - ctx_start + num_gen_reqs = gen_end - gen_start + + if self.is_cuda_graph: + # CUDA graph capture bakes host-side launch scalars into + # thop.attention. Use conservative per-domain bounds so a graph + # captured for one decode length can replay for later, longer + # decode steps with the same request split. + ctx_total_kv_length = self.max_seq_len * num_ctx_reqs + gen_total_kv_length = self.max_seq_len * num_gen_reqs + else: + ctx_total_kv_length = ( + kv_lens[ctx_start:ctx_end].sum().item() + if num_ctx_reqs > 0 else 0) + gen_total_kv_length = ( + kv_lens[gen_start:gen_end].sum().item() + if num_gen_reqs > 0 else 0) + buf = self._locality_domain_host_total_kv_lens_buf[ + locality_domain_idx] + buf[0] = ctx_total_kv_length + buf[1] = gen_total_kv_length + + self.locality_domain_mla_gen_state = [None + ] * num_locality_domains + # Don't pass self.kv_lens as kv_lens here because it includes extra # tokens. Use the actual KV length (without extra tokens) for # kv_lens_runtime, which becomes host_past_key_value_lengths and @@ -1139,6 +1389,12 @@ def prepare(self) -> None: kv_lens_runtime = kv_lens[:self.num_seqs] prompt_lens_cuda_runtime = self.prompt_lens_cuda[:self.num_seqs] prompt_lens_cpu_runtime = self.prompt_lens_cpu[:self.num_seqs] + # Under locality-domain CUDA graphs, host-side launch scalars are captured per + # sub-call, so use conservative bounds while keeping the device sequence + # tensors and prompt lengths actual. + if self.is_cuda_graph and self.locality_domain_enabled: + kv_lens_runtime = torch.full_like(kv_lens[:self.num_seqs], + self.max_seq_len) self._bind_runtime_views( kv_lens_cuda=self.kv_lens_cuda[:self.num_seqs], kv_lens=kv_lens_runtime, @@ -1527,6 +1783,21 @@ def spec_decoding_param_prepare_for_blackwell(self) -> None: device='cuda', ) + def _ensure_locality_domain_spec_bl_tree_masks(self) -> None: + if not self.locality_domain_enabled or self.kv_cache_manager is None: + return + num_locality_domains = self.kv_cache_manager.num_locality_domains + if num_locality_domains <= 1 or self.spec_decoding_bl_tree_mask is None: + return + masks = self.locality_domain_spec_decoding_bl_tree_masks + if (masks is None or len(masks) != num_locality_domains + or any(mask.shape != self.spec_decoding_bl_tree_mask.shape + for mask in masks[1:])): + self.locality_domain_spec_decoding_bl_tree_masks = [None] + [ + torch.zeros_like(self.spec_decoding_bl_tree_mask) + for _ in range(num_locality_domains - 1) + ] + def update_blackwell_first_sparse_mask_offset(self) -> None: """Fill gen slots [0:ng) with (kv_lens_cuda - seq_lens); in-place for CUDA graph.""" nc = self.num_contexts @@ -1623,10 +1894,12 @@ def update_spec_dec_param( if self.is_sm_version_trtllm_gen_kernel(sm=get_sm_version()): self.spec_decoding_param_prepare_for_blackwell() + self._ensure_locality_domain_spec_bl_tree_masks() else: self.spec_decoding_bl_tree_mask_offset = None self.spec_decoding_bl_tree_mask = None self.spec_bl_tree_first_sparse_mask_offset_kv = None + self.locality_domain_spec_decoding_bl_tree_masks = None cpp_query_len = 0 @@ -1701,6 +1974,635 @@ def is_sm_version_trtllm_gen_kernel(self, sm): return not (sm < 100 or sm in [120, 121]) +def _get_generation_only_token_splits( + metadata: TrtllmAttentionMetadata, + total_gen_tokens: Optional[int], +) -> Tuple[List[int], List[int]]: + gen_token_splits = metadata.locality_domain_gen_token_splits + per_locality_domain_gen_tokens = metadata.locality_domain_num_gen_tokens + + if (total_gen_tokens is None or gen_token_splits is None + or per_locality_domain_gen_tokens is None + or metadata.num_generations == 0): + return gen_token_splits, per_locality_domain_gen_tokens + + expected_total_gen_tokens = gen_token_splits[-1] + if total_gen_tokens == expected_total_gen_tokens: + return gen_token_splits, per_locality_domain_gen_tokens + + # MTP draft forwards can shrink the live q/fused_q tensor without + # rebuilding the cached locality domain token splits. Request ordering is still + # locality domain-contiguous, so rebuild generation token boundaries from the + # current token count for this forward. + if total_gen_tokens % metadata.num_generations != 0: + raise RuntimeError( + "locality domain generation-only attention received a token count that " + "cannot be partitioned evenly across generation requests: " + f"{total_gen_tokens=} {metadata.num_generations=}") + + tokens_per_generation_request = total_gen_tokens // metadata.num_generations + current_gen_token_splits = [0] + current_per_locality_domain_gen_tokens = [] + for locality_domain_idx in range( + len(metadata.locality_domain_gen_req_splits) - 1): + num_gen_reqs = ( + metadata.locality_domain_gen_req_splits[locality_domain_idx + 1] - + metadata.locality_domain_gen_req_splits[locality_domain_idx]) + locality_domain_gen_tokens = num_gen_reqs * tokens_per_generation_request + current_per_locality_domain_gen_tokens.append( + locality_domain_gen_tokens) + current_gen_token_splits.append(current_gen_token_splits[-1] + + locality_domain_gen_tokens) + + return current_gen_token_splits, current_per_locality_domain_gen_tokens + + +def _build_locality_domain_slices(metadata: TrtllmAttentionMetadata, + locality_domain_id: int, + attention_input_type: AttentionInputType, + total_gen_tokens: Optional[int] = None, + dense_ctx_kv: bool = False): + """Build per-locality domain request/token slices and metadata tensors. + + Returns a dict with per-locality domain counts, slices, and metadata tensors, + or None if this locality domain has no tokens for the given phase. + + When ``dense_ctx_kv`` is True (MLA cached-KV / chunked-prefill context), + ctx k/v rows are laid out by per-request KV length (cached+new or + chunk-sized), not by per-request Q length. The returned + ``ctx_kv_token_start/end`` give the per-locality domain k/v slice for that layout; + callers continue to use ``ctx_token_start/end`` for q/q_pe. When False, + the KV range equals the Q range (current behavior). + """ + num_ctx = metadata.num_contexts + ctx_req_splits = metadata.locality_domain_ctx_req_splits + gen_req_splits = metadata.locality_domain_gen_req_splits + ctx_token_splits = metadata.locality_domain_ctx_token_splits + gen_token_splits = metadata.locality_domain_gen_token_splits + total_ctx_tokens = sum(metadata.locality_domain_num_ctx_tokens) + if attention_input_type == AttentionInputType.generation_only: + gen_token_splits, _ = _get_generation_only_token_splits( + metadata, total_gen_tokens) + + # Request slices via cumulative boundary arrays + ctx_req_start = ctx_req_splits[locality_domain_id] + ctx_req_end = ctx_req_splits[locality_domain_id + 1] + gen_req_start = num_ctx + gen_req_splits[locality_domain_id] + gen_req_end = num_ctx + gen_req_splits[locality_domain_id + 1] + + # Token slices depend on attention_input_type + if attention_input_type == AttentionInputType.generation_only: + ctx_token_start, ctx_token_end = 0, 0 + gen_token_start = gen_token_splits[locality_domain_id] + gen_token_end = gen_token_splits[locality_domain_id + 1] + elif attention_input_type == AttentionInputType.context_only: + # context_only: q/k/v tensors have ONLY context tokens. + # Gen token slices must be zero so _locality_domain_cat_slices returns a + # view (preserving strides) instead of torch.cat (which makes + # a contiguous copy). A contiguous V copy breaks the TMA + # descriptor stride assumption in the FMHA kernel. + ctx_token_start = ctx_token_splits[locality_domain_id] + ctx_token_end = ctx_token_splits[locality_domain_id + 1] + gen_token_start, gen_token_end = 0, 0 + else: + # mixed: combined [ctx_tokens | gen_tokens] layout + ctx_token_start = ctx_token_splits[locality_domain_id] + ctx_token_end = ctx_token_splits[locality_domain_id + 1] + gen_token_start = total_ctx_tokens + gen_token_splits[locality_domain_id] + gen_token_end = total_ctx_tokens + gen_token_splits[locality_domain_id + + 1] + + # For generation_only (MLA), exclude context requests from per-request + # metadata. The mla_rope_generation C++ kernel expects token tensors + # and per-request metadata to be consistent — token buffers have only + # gen tokens, so the metadata must also cover only gen requests. + if attention_input_type == AttentionInputType.generation_only: + ctx_req_start = ctx_req_end # zero-width context range + + num_ctx_reqs = ctx_req_end - ctx_req_start + num_gen_reqs = gen_req_end - gen_req_start + num_ctx_tokens = ctx_token_end - ctx_token_start + num_gen_tokens = gen_token_end - gen_token_start + num_tokens = num_ctx_tokens + num_gen_tokens + + # Dense-KV ctx slicing: rows of k/v are laid out by per-req kv_len + # (cached+new or chunk-sized), contiguous per request. The per-locality domain + # boundary is the running sum of ``kv_lens_runtime`` over preceding + # locality domains' ctx requests. Derived dynamically so chunked prefill (which + # rewrites ``kv_lens_runtime`` per loop iter) works without extra + # metadata. When not dense, fall through to the Q range. + if dense_ctx_kv and attention_input_type == AttentionInputType.context_only: + assert metadata.kv_lens_runtime is not None, ( + "dense_ctx_kv requires kv_lens_runtime to be prepared") + all_ctx_kv_lens = metadata.kv_lens_runtime[:num_ctx] + ctx_kv_token_start = int( + all_ctx_kv_lens[:ctx_req_splits[locality_domain_id]].sum().item()) + ctx_kv_token_end = int( + all_ctx_kv_lens[:ctx_req_splits[locality_domain_id + + 1]].sum().item()) + else: + ctx_kv_token_start = ctx_token_start + ctx_kv_token_end = ctx_token_end + + if num_tokens == 0: + return None + if attention_input_type == AttentionInputType.context_only and num_ctx_tokens == 0: + return None + if attention_input_type == AttentionInputType.generation_only and num_gen_tokens == 0: + return None + + # Per-request metadata: repack ctx + gen contiguously + num_seqs = num_ctx_reqs + num_gen_reqs + if num_seqs == 0: + return None + if attention_input_type == AttentionInputType.generation_only: + # Gen-only: context range is empty (ctx_req_start==ctx_req_end), + # use views directly. + # Eliminates torch.cat allocations — CUDA-graph-safe. + sequence_lengths = metadata.kv_lens_cuda_runtime[ + gen_req_start:gen_req_end] + kv_lengths = metadata.kv_lens_runtime[gen_req_start:gen_req_end] + prompt_lengths_cuda = metadata.prompt_lens_cuda_runtime[ + gen_req_start:gen_req_end] + prompt_lengths_cpu = metadata.prompt_lens_cpu_runtime[ + gen_req_start:gen_req_end] + host_context_lengths = prompt_lengths_cpu + if metadata.is_cuda_graph: + # attentionOp.cpp captures max(host_context_lengths) as a launch + # scalar. Use the same graph launch bounds as host KV lengths for + # gen-only locality domain graphs so a graph captured on a short prompt can + # replay on later, longer prompts. + host_context_lengths = kv_lengths + request_types = metadata.host_request_types_runtime[ + gen_req_start:gen_req_end] + host_total_kv_lengths = metadata._locality_domain_host_total_kv_lens_buf[ + locality_domain_id] + else: + sequence_lengths = torch.cat([ + metadata.kv_lens_cuda_runtime[ctx_req_start:ctx_req_end], + metadata.kv_lens_cuda_runtime[gen_req_start:gen_req_end] + ]) + kv_lengths = torch.cat([ + metadata.kv_lens_runtime[ctx_req_start:ctx_req_end], + metadata.kv_lens_runtime[gen_req_start:gen_req_end] + ]) + prompt_lengths_cuda = torch.cat([ + metadata.prompt_lens_cuda_runtime[ctx_req_start:ctx_req_end], + metadata.prompt_lens_cuda_runtime[gen_req_start:gen_req_end] + ]) + prompt_lengths_cpu = torch.cat([ + metadata.prompt_lens_cpu_runtime[ctx_req_start:ctx_req_end], + metadata.prompt_lens_cpu_runtime[gen_req_start:gen_req_end] + ]) + host_context_lengths = prompt_lengths_cpu + request_types = torch.cat([ + metadata.host_request_types_runtime[ctx_req_start:ctx_req_end], + metadata.host_request_types_runtime[gen_req_start:gen_req_end] + ]) + if dense_ctx_kv: + # Dense-KV callers (MLA cached-KV / chunked-prefill context) rebind + # ``metadata.kv_lens_runtime`` per call to reflect the actual row + # count of the passed k/v, so the stale + # ``_locality_domain_host_total_kv_lens_buf`` captured at prepare() no longer + # matches the kernel inputs. Recompute from the live runtime tensor. + host_total_kv_lengths = torch.tensor([ + int(metadata.kv_lens_runtime[ctx_req_start:ctx_req_end].sum(). + item()), + int(metadata.kv_lens_runtime[gen_req_start:gen_req_end].sum(). + item()) + ], + dtype=torch.int, + device='cpu') + else: + host_total_kv_lengths = metadata._locality_domain_host_total_kv_lens_buf[ + locality_domain_id] + + # The KV cache manager is shared across locality domains; per-locality domain addressing + # is encoded in the block offsets (KVCacheManagerV2 uses a single + # canonical pool pointer, with the locality domain-specific virtual-memory + # offset baked into the global page indices). + kv_cache_manager = metadata.kv_cache_manager + + # The per-locality domain block offsets tensor is laid out as + # [ctx_block_offsets | gen_block_offsets]. For generation_only, the + # C++ kernel uses num_ctx_reqs=0 and expects gen entries at position 0. + # Re-index the block offsets so gen entries start at position 0. + block_offsets = metadata.locality_domain_kv_cache_block_offsets[ + locality_domain_id] + if attention_input_type == AttentionInputType.generation_only and num_gen_reqs > 0: + # Compute the ctx count in the block offsets layout (from + # prepare_batch, which packs ctx requests before gen requests). + num_ctx_in_block_offsets = (ctx_req_splits[locality_domain_id + 1] - + ctx_req_splits[locality_domain_id]) + if num_ctx_in_block_offsets > 0: + gen_block_offsets = torch.zeros_like(block_offsets) + gen_block_offsets[:, :num_gen_reqs].copy_( + block_offsets[:, num_ctx_in_block_offsets: + num_ctx_in_block_offsets + num_gen_reqs]) + block_offsets = gen_block_offsets + + return { + 'ctx_req_start': ctx_req_start, + 'ctx_req_end': ctx_req_end, + 'gen_req_start': gen_req_start, + 'gen_req_end': gen_req_end, + 'ctx_token_start': ctx_token_start, + 'ctx_token_end': ctx_token_end, + 'ctx_kv_token_start': ctx_kv_token_start, + 'ctx_kv_token_end': ctx_kv_token_end, + 'gen_token_start': gen_token_start, + 'gen_token_end': gen_token_end, + 'num_ctx_reqs': num_ctx_reqs, + 'num_gen_reqs': num_gen_reqs, + 'num_ctx_tokens': num_ctx_tokens, + 'num_gen_tokens': num_gen_tokens, + 'num_tokens': num_tokens, + 'num_seqs': num_seqs, + 'sequence_lengths': sequence_lengths, + 'kv_lengths': kv_lengths, + 'prompt_lengths_cuda': prompt_lengths_cuda, + 'prompt_lengths_cpu': prompt_lengths_cpu, + 'host_context_lengths': host_context_lengths, + 'request_types': request_types, + 'host_total_kv_lengths': host_total_kv_lengths, + 'block_offsets': block_offsets, + 'pool_pointers': kv_cache_manager.kv_cache_pool_pointers, + 'pool_mapping': metadata.host_kv_cache_pool_mapping, + 'kv_cache_manager': kv_cache_manager, + } + + +def _check_locality_domain_side_inputs(metadata, forward_args) -> None: + """Reject per-request inputs the partition path does not slice. + + Everything listed here is indexed by batch row or by token inside the + kernels, so handing a domain the full-batch tensor would read another + domain's rows. Raising beats a silent full-batch fallback that hides the + missing locality. + """ + unsupported = [] + # The scoped view clears block_ids_per_seq, which is the FlashMLA block table. + if getattr(metadata, "enable_flash_mla", False): + unsupported.append("FlashMLA") + if metadata.beam_width != 1: + unsupported.append(f"beam_width={metadata.beam_width}") + if metadata.is_cross or forward_args.cross_kv is not None: + unsupported.append("cross attention") + if metadata.enable_helix: + unsupported.append("helix parallelism") + if forward_args.attention_mask_data is not None: + unsupported.append("custom attention mask data") + sparse = forward_args.sparse_runtime_params + if sparse is not None and sparse.sparse_attn_indices is not None: + unsupported.append("sparse attention indices") + if (forward_args.mrope_rotary_cos_sin is not None + and forward_args.attention_input_type != AttentionInputType.mixed): + unsupported.append("MRoPE outside the mixed attention input type") + if unsupported: + raise NotImplementedError( + "Locality domain fork-join attention does not support " + f"{', '.join(unsupported)}; disable enable_locality_domains for " + "this configuration.") + + +def _locality_domain_attention_enabled(attn, metadata, forward_args) -> bool: + """Whether this call should be split across locality domains.""" + # The flag is the explicit opt-out: draft sub-forwards without their own + # per-domain layout clear it while the manager still reports two domains. + if not metadata.locality_domain_enabled: + return False + # output_sf is deliberately not checked here: an NVFP4 output quantization + # request that reaches this path is asserted on rather than silently run + # unpartitioned. + if forward_args.output is None: + return False + # MLA chunked-prefill context gathers k/v by per-request KV length, and a + # domain can hold a mix of chunk lengths including zero, which FMHA rejects. + # Leave it on the single-kernel path; generation in the same step still splits. + if (attn.is_mla_enable and forward_args.attention_input_type + == AttentionInputType.context_only + and forward_args.latent_cache is None + and metadata.runtime_features is not None + and metadata.runtime_features.chunked_prefill): + return False + kv_cache_manager = metadata.kv_cache_manager + if kv_cache_manager is None: + return False + if getattr(kv_cache_manager, "num_locality_domains", 1) <= 1: + return False + if metadata.locality_domain_gen_req_splits is None: + return False + _check_locality_domain_side_inputs(metadata, forward_args) + return True + + +def _resize_locality_domain_workspaces_from_domain0(metadata) -> None: + """Propagate a workspace growth on domain 0 to the other domains. + + The attention kernel resizes the workspace it is handed, so with one + workspace per domain only domain 0's grows. + """ + workspaces = metadata.locality_domain_workspaces + if metadata.is_cuda_graph or workspaces is None: + return + target_numel = workspaces[0].numel() + if target_numel == 0: + return + runtime = metadata.locality_domain_runtime + for locality_domain_id, workspace in enumerate(workspaces[1:], start=1): + if workspace.numel() >= target_numel: + continue + with runtime.partition_context(locality_domain_id): + workspace.resize_(target_numel) + + +_LOCALITY_DOMAIN_GEN_INDEXED_SPEC_FIELDS = ( + "spec_decoding_generation_lengths", + "spec_decoding_packed_mask", + "spec_decoding_bl_tree_mask_offset", + "spec_bl_tree_first_sparse_mask_offset_kv", +) + + +def _locality_domain_mrope_overrides(forward_args, ld_slice, num_contexts, + num_q_rows): + """Slice the MRoPE inputs the way the kernels index them. + + The context kernel reads cos/sin by batch-flat token row and the generation + kernels read position deltas by local generation batch index. + """ + overrides = {} + cos_sin = forward_args.mrope_rotary_cos_sin + if cos_sin is not None: + ctx_start, ctx_end = ld_slice["ctx_token_start"], ld_slice[ + "ctx_token_end"] + if ctx_end > ctx_start: + if num_q_rows <= 0 or cos_sin.numel() % num_q_rows != 0: + raise RuntimeError( + f"mrope_rotary_cos_sin holds {cos_sin.numel()} values for " + f"{num_q_rows} tokens; cannot slice it per locality domain") + per_token = cos_sin.numel() // num_q_rows + overrides["mrope_rotary_cos_sin"] = cos_sin.reshape( + -1)[ctx_start * per_token:ctx_end * per_token].view(1, -1) + else: + overrides["mrope_rotary_cos_sin"] = None + deltas = forward_args.mrope_position_deltas + if deltas is not None: + gen_start = ld_slice["gen_req_start"] - num_contexts + gen_end = ld_slice["gen_req_end"] - num_contexts + overrides["mrope_position_deltas"] = (deltas[gen_start:gen_end] + if gen_end > gen_start else None) + return overrides + + +def _spans_tile(intervals, start, end) -> bool: + if any(a < 0 or b < a for a, b in intervals): + return False + spans = sorted((a, b) for a, b in intervals if b > a) + if not spans: + return start == end + if spans[0][0] != start or spans[-1][1] != end: + return False + return all(spans[i][1] == spans[i + 1][0] for i in range(len(spans) - 1)) + + +def _validate_locality_domain_coverage(metadata, prepared, attention_input_type, + num_q_rows, num_kv_rows, + dense_ctx_kv) -> None: + """Check on the host that the domain slices tile the live batch exactly. + + Splits are built at prepare(); a draft rewrite that changes the layout + without a new prepare() would otherwise launch kernels over a subset of + the rows and leave the rest of the output unwritten. Python ints only, so + this never synchronizes. + """ + slices = [prep["ld_slice"] for prep in prepared] + + def spans(start_key, end_key): + return [(s[start_key], s[end_key]) for s in slices] + + num_ctx = metadata.num_contexts + num_gen = metadata.num_generations + problems = [] + batch_rows = metadata.kv_lens_runtime.shape[0] + if batch_rows != num_ctx + num_gen: + problems.append(f"runtime metadata has {batch_rows} rows for " + f"{num_ctx} context + {num_gen} generation requests") + for s in slices: + if (s["num_ctx_tokens"] != s["ctx_token_end"] - s["ctx_token_start"] + or s["num_gen_tokens"] + != s["gen_token_end"] - s["gen_token_start"]): + problems.append(f"inconsistent token counts in slice {s}") + ctx_tokens = spans("ctx_token_start", "ctx_token_end") + gen_tokens = spans("gen_token_start", "gen_token_end") + ctx_reqs = spans("ctx_req_start", "ctx_req_end") + gen_reqs = spans("gen_req_start", "gen_req_end") + if attention_input_type == AttentionInputType.generation_only: + expected = [("generation tokens", gen_tokens, 0, num_q_rows), + ("generation requests", gen_reqs, num_ctx, + num_ctx + num_gen), ("context tokens", ctx_tokens, 0, 0)] + elif attention_input_type == AttentionInputType.context_only: + expected = [("context tokens", ctx_tokens, 0, num_q_rows), + ("context requests", ctx_reqs, 0, num_ctx), + ("generation tokens", gen_tokens, 0, 0)] + ctx_kv = spans("ctx_kv_token_start", "ctx_kv_token_end") + if dense_ctx_kv: + if num_kv_rows is not None: + expected.append(("context kv rows", ctx_kv, 0, num_kv_rows)) + elif ctx_kv != ctx_tokens: + problems.append(f"context kv rows {ctx_kv} differ from context " + f"tokens {ctx_tokens}") + else: + total_ctx_tokens = sum(metadata.locality_domain_num_ctx_tokens) + expected = [ + ("context tokens", ctx_tokens, 0, total_ctx_tokens), + ("generation tokens", gen_tokens, total_ctx_tokens, num_q_rows), + ("context requests", ctx_reqs, 0, num_ctx), + ("generation requests", gen_reqs, num_ctx, num_ctx + num_gen) + ] + for label, intervals, start, end in expected: + if not _spans_tile(intervals, start, end): + problems.append(f"{label} {intervals} do not tile [{start}, {end})") + if problems: + raise RuntimeError( + "Locality domain attention slices do not cover the batch " + f"({attention_input_type.name}): " + "; ".join(problems)) + + +def _scatter_locality_domain_rows(dst, src, ld_slice, num_ctx_tokens, + num_gen_tokens) -> None: + if num_ctx_tokens > 0: + dst[ld_slice["ctx_token_start"]:ld_slice["ctx_token_end"]] = ( + src[:num_ctx_tokens]) + if num_gen_tokens > 0: + dst[ld_slice["gen_token_start"]:ld_slice["gen_token_end"]] = ( + src[num_ctx_tokens:]) + + +@contextmanager +def _locality_domain_metadata(metadata, ld_slice, locality_domain_id): + """Point the shared metadata at one locality domain's rows for the duration. + + The FMHA backends read their per-request tensors off ``metadata`` rather than + taking them as arguments, so a per-domain call is expressed by swapping those + fields for the domain's slices and restoring them afterwards. + """ + overrides = { + "kv_lens_cuda_runtime": ld_slice["sequence_lengths"], + "kv_lens_runtime": ld_slice["kv_lengths"], + "prompt_lens_cuda_runtime": ld_slice["prompt_lengths_cuda"], + "prompt_lens_cpu_runtime": ld_slice["prompt_lengths_cpu"], + "host_request_types_runtime": ld_slice["request_types"], + "host_total_kv_lens": ld_slice["host_total_kv_lengths"], + "kv_cache_block_offsets": ld_slice["block_offsets"], + # pool pointers/mapping are read-only and identical for every domain. + # The counts go on the private fields: the property setters call + # on_update(), which would recompute them from the full-batch seq_lens. + "_num_contexts": ld_slice["num_ctx_reqs"], + "_num_generations": ld_slice["num_gen_reqs"], + "_num_ctx_tokens": ld_slice["num_ctx_tokens"], + "_num_tokens": ld_slice["num_tokens"], + # Indexed by sequence over the whole batch, so it cannot describe a + # single domain's rows. + "block_ids_per_seq": None, + } + if metadata.use_spec_decoding: + # The generation kernels index these by the local generation batch + # index, so each domain gets its own rows. + gen_start = ld_slice["gen_req_start"] - metadata.num_contexts + gen_end = ld_slice["gen_req_end"] - metadata.num_contexts + for name in _LOCALITY_DOMAIN_GEN_INDEXED_SPEC_FIELDS: + tensor = getattr(metadata, name, None) + if tensor is not None: + overrides[name] = tensor[gen_start:gen_end] + if metadata.spec_decoding_position_offsets is not None: + # Dynamic trees write the offsets compactly, row width + # max_total_draft_tokens + 1, into a buffer whose allocation width + # may be larger; slice the compact view so later domains start at + # the right row. + offsets = metadata.spec_decoding_position_offsets_cpp + if offsets is None or offsets.dim() != 2: + offsets = metadata.spec_decoding_position_offsets_for_cpp + offsets = offsets[gen_start:gen_end] + overrides["spec_decoding_position_offsets"] = offsets + overrides["spec_decoding_position_offsets_cpp"] = offsets + # prepareCustomMask writes the trtllm-gen tree mask from offset 0 on + # every call, so concurrent domains cannot share one mask buffer. + masks = metadata.locality_domain_spec_decoding_bl_tree_masks + if (masks is not None and locality_domain_id < len(masks) + and masks[locality_domain_id] is not None): + overrides["spec_decoding_bl_tree_mask"] = masks[locality_domain_id] + # Each domain needs its own scratch: the streams run concurrently, so a + # shared workspace would be a data race. ``effective_workspace`` picks + # between these two by ``is_cuda_graph``, so override both. + for name, per_domain in (("workspace", metadata.locality_domain_workspaces), + ("cuda_graph_workspace", + metadata.locality_domain_cuda_graph_workspaces)): + if per_domain is not None and locality_domain_id < len(per_domain): + overrides[name] = per_domain[locality_domain_id] + saved = {k: getattr(metadata, k) for k in overrides} + try: + for name, value in overrides.items(): + setattr(metadata, name, value) + yield + finally: + for name, value in saved.items(): + setattr(metadata, name, value) + + +def _locality_domain_cat_slices(t, ctx_token_start, ctx_token_end, + gen_token_start, gen_token_end, num_ctx_tokens, + num_gen_tokens): + """Cat context and generation token slices for a locality domain.""" + if t is None: + return None + if num_ctx_tokens > 0 and num_gen_tokens > 0: + return torch.cat([ + t[ctx_token_start:ctx_token_end], t[gen_token_start:gen_token_end] + ], + dim=0) + return t[ctx_token_start:ctx_token_end] if num_gen_tokens == 0 else t[ + gen_token_start:gen_token_end] + + +def _slice_cu_seqlens(cu_seqlens: torch.Tensor, start: int, + end: int) -> torch.Tensor: + sliced = cu_seqlens[start:end + 1] - cu_seqlens[start] + return sliced.contiguous() + + +def _slice_locality_domain_cu_seqlens(cu_seqlens, ld_slice, + attention_input_type, + num_contexts) -> torch.Tensor: + """Per-domain view of a full-batch cumulative length tensor. + + The full tensor is indexed by request in batch order. Generation-only + calls carry generation requests only, so their indices drop the context + offset; mixed calls splice the domain's context and generation ranges. + """ + if attention_input_type == AttentionInputType.generation_only: + return _slice_cu_seqlens(cu_seqlens, + ld_slice["gen_req_start"] - num_contexts, + ld_slice["gen_req_end"] - num_contexts) + if attention_input_type == AttentionInputType.context_only: + return _slice_cu_seqlens(cu_seqlens, ld_slice["ctx_req_start"], + ld_slice["ctx_req_end"]) + ctx = _slice_cu_seqlens(cu_seqlens, ld_slice["ctx_req_start"], + ld_slice["ctx_req_end"]) + gen = _slice_cu_seqlens(cu_seqlens, ld_slice["gen_req_start"], + ld_slice["gen_req_end"]) + return torch.cat((ctx, gen[1:] + ctx[-1])) + + +def _prepare_locality_domain_mla_gen_state_from_full( + metadata: TrtllmAttentionMetadata, + cu_q_seqlens: torch.Tensor, + cu_kv_seqlens: torch.Tensor, + fmha_scheduler_counter: torch.Tensor, + mla_bmm1_scale: Optional[torch.Tensor], + mla_bmm2_scale: Optional[torch.Tensor], + total_gen_tokens: int, +) -> None: + num_locality_domains = metadata.kv_cache_manager.num_locality_domains + metadata.locality_domain_mla_gen_state = [None] * num_locality_domains + counters = ( + fmha_scheduler_counter.new_zeros((num_locality_domains, ) + + tuple(fmha_scheduler_counter.shape)) + if fmha_scheduler_counter is not None else None) + for locality_domain_id in range(num_locality_domains): + locality_domain_slice = _build_locality_domain_slices( + metadata, + locality_domain_id, + AttentionInputType.generation_only, + total_gen_tokens=total_gen_tokens, + ) + if locality_domain_slice is None: + continue + + gen_req_start = locality_domain_slice[ + 'gen_req_start'] - metadata.num_contexts + gen_req_end = locality_domain_slice[ + 'gen_req_end'] - metadata.num_contexts + counter = counters[locality_domain_id] if counters is not None else None + # The rope kernel already wrote these and FMHA only reads them, and + # cumulative lengths start at 0, so a slice from row 0 can be a view. + if gen_req_start == 0: + cu_q_seqlens_locality_domain = cu_q_seqlens[:gen_req_end + 1] + cu_kv_seqlens_locality_domain = cu_kv_seqlens[:gen_req_end + 1] + else: + cu_q_seqlens_locality_domain = _slice_cu_seqlens( + cu_q_seqlens, gen_req_start, gen_req_end) + cu_kv_seqlens_locality_domain = _slice_cu_seqlens( + cu_kv_seqlens, gen_req_start, gen_req_end) + metadata.locality_domain_mla_gen_state[locality_domain_id] = { + 'cu_q_seqlens': cu_q_seqlens_locality_domain, + 'cu_kv_seqlens': cu_kv_seqlens_locality_domain, + 'fmha_scheduler_counter': counter, + 'mla_bmm1_scale': mla_bmm1_scale, + 'mla_bmm2_scale': mla_bmm2_scale, + } + + class TrtllmAttention(AttentionBackend[TrtllmAttentionMetadata]): Metadata = TrtllmAttentionMetadata @@ -2505,6 +3407,11 @@ def forward( self.local_layer_idx = self.get_local_layer_idx(metadata) if metadata.spec_decoding_bl_tree_mask is not None and self.local_layer_idx == 0: metadata.spec_decoding_bl_tree_mask.zero_() + # prepareCustomMask ORs into whichever mask a domain is handed. + domain_masks = metadata.locality_domain_spec_decoding_bl_tree_masks + for mask in domain_masks or (): + if mask is not None: + mask.zero_() if self.print_skip_softmax_stat: self.skip_softmax_stat.zero_() @@ -2560,29 +3467,32 @@ def forward( # library selection and dispatch: support checks read the # same runtime metadata as the selected FMHA implementation. with metadata.presented_token_major(token_major_view): - fmha = self._fmha_manager.select(self, q, k, v, metadata, - forward_args) - - if fmha is None: - if (getattr(forward_args, "qk_fp8_per_tensor", False) - or getattr(forward_args, "v_fp8_per_tensor", False)): + for q_part, k_part, v_part, part_args in self._attention_partitions( + q, k, v, metadata, forward_args): + fmha = self._fmha_manager.select(self, q_part, k_part, v_part, + metadata, part_args) + + if fmha is None: + if (getattr(part_args, "qk_fp8_per_tensor", False) + or getattr(part_args, "v_fp8_per_tensor", False)): + raise RuntimeError( + "The per-tensor FP8 attention recipe needs the PrimTS " + "FMHA library (TLLM_FMHA_LIBS=+prims_ts), which is not " + "enabled or does not support this request.") raise RuntimeError( - "The per-tensor FP8 attention recipe needs the PrimTS " - "FMHA library (TLLM_FMHA_LIBS=+prims_ts), which is not " - "enabled or does not support this request.") - raise RuntimeError( - "No TRT-LLM attention FMHA library supports this request.") - if metadata.is_cuda_graph or not fmha.supports_workspace_reclamation: - # Conservatively disable reclamation for metadata used by graphs - # or backends that can retain staged workspace state. - metadata.workspace_reclaimable = False - try: - fmha.forward(q, k, v, metadata, forward_args) - except RuntimeError as exc: - log_attention_failure_context( - type(self).__name__, self.layer_idx, metadata, - forward_args.attention_window_size, exc) - raise + "No TRT-LLM attention FMHA library supports this request." + ) + if metadata.is_cuda_graph or not fmha.supports_workspace_reclamation: + # Conservatively disable reclamation for metadata used by graphs + # or backends that can retain staged workspace state. + metadata.workspace_reclaimable = False + try: + fmha.forward(q_part, k_part, v_part, metadata, part_args) + except RuntimeError as exc: + log_attention_failure_context( + type(self).__name__, self.layer_idx, metadata, + part_args.attention_window_size, exc) + raise if self.print_skip_softmax_stat: total_blocks, skipped_blocks = self.skip_softmax_stat @@ -2596,6 +3506,197 @@ def forward( else: return forward_args.output, forward_args.output_sf + def _attention_partitions(self, q, k, v, metadata, forward_args): + """Yield the (q, k, v, forward_args) partitions this call must dispatch. + + Without locality domains this yields the batch once, unchanged. With them + it yields once per domain -- inputs sliced to that domain's rows, its own + workspace and MLA state in place, and its stream current -- so the caller + keeps a single dispatch site for both paths. + + Generation rows are split by position (the runner's balanced half/half + split), not by the domain that owns their KV pages, so a row can run on + the other domain's stream; the result is still correct. Context rows + follow their owning domain when the runner could sort them. + """ + if not _locality_domain_attention_enabled(self, metadata, forward_args): + yield q, k, v, forward_args + return + + assert forward_args.output_sf is None, ( + "NVFP4 output quantization is not supported with locality domain " + "fork-join attention.") + + num_locality_domains = metadata.kv_cache_manager.num_locality_domains + # MLA context with pre-gathered k/v lays those rows out by per-request + # KV length, not by Q length, so the k/v slice differs from q's. + is_gen_only = (forward_args.attention_input_type == + AttentionInputType.generation_only) + dense_ctx_kv = (self.is_mla_enable and forward_args.attention_input_type + == AttentionInputType.context_only + and forward_args.latent_cache is None) + if metadata.locality_domain_runtime is None: + metadata.locality_domain_runtime = LocalityDomainRuntime( + num_partitions=num_locality_domains) + runtime = metadata.locality_domain_runtime + output = forward_args.output + + # Prep every domain first, on the default stream: no host-side tensor + # building happens inside the fork/join region. + prepared = [] + for locality_domain_id in range(num_locality_domains): + ld_slice = _build_locality_domain_slices( + metadata, + locality_domain_id, + forward_args.attention_input_type, + total_gen_tokens=(q.shape[0] if is_gen_only else None), + dense_ctx_kv=dense_ctx_kv) + if ld_slice is None: + continue + num_ctx_tokens = ld_slice["num_ctx_tokens"] + num_gen_tokens = ld_slice["num_gen_tokens"] + if num_ctx_tokens + num_gen_tokens == 0: + continue + num_ctx_kv_tokens = (ld_slice["ctx_kv_token_end"] - + ld_slice["ctx_kv_token_start"]) + cat = functools.partial(_locality_domain_cat_slices, + ctx_token_start=ld_slice["ctx_token_start"], + ctx_token_end=ld_slice["ctx_token_end"], + gen_token_start=ld_slice["gen_token_start"], + gen_token_end=ld_slice["gen_token_end"], + num_ctx_tokens=num_ctx_tokens, + num_gen_tokens=num_gen_tokens) + cat_kv = functools.partial( + _locality_domain_cat_slices, + ctx_token_start=ld_slice["ctx_kv_token_start"], + ctx_token_end=ld_slice["ctx_kv_token_end"], + gen_token_start=ld_slice["gen_token_start"], + gen_token_end=ld_slice["gen_token_end"], + num_ctx_tokens=num_ctx_kv_tokens, + num_gen_tokens=num_gen_tokens) + if num_ctx_tokens > 0 and num_gen_tokens > 0: + ld_output = output.new_empty((num_ctx_tokens + + num_gen_tokens, ) + + output.shape[1:]) + scatter_back = True + else: + ld_output = cat(output) + scatter_back = False + ld_overrides = dict(output=ld_output, + locality_domain_id=locality_domain_id) + # Written per token row by the kernel, so it needs the same + # view-or-copy treatment as the output. + softmax_stats = forward_args.softmax_stats_tensor + ld_softmax_stats = None + if softmax_stats is not None: + if scatter_back: + ld_softmax_stats = softmax_stats.new_empty( + (num_ctx_tokens + num_gen_tokens, ) + + softmax_stats.shape[1:]) + else: + ld_softmax_stats = cat(softmax_stats) + ld_overrides["softmax_stats_tensor"] = ld_softmax_stats + ld_overrides.update( + _locality_domain_mrope_overrides(forward_args, ld_slice, + metadata.num_contexts, + q.shape[0])) + for name in ("latent_cache", "q_pe"): + tensor = getattr(forward_args, name) + if tensor is not None: + ld_overrides[name] = cat(tensor) + mla_state = metadata.locality_domain_mla_gen_state + mla_state = (mla_state[locality_domain_id] if + (mla_state is not None + and locality_domain_id < len(mla_state)) else None) + if mla_state is not None and is_gen_only: + ld_overrides.update(mla_state) + if forward_args.quant_q_buffer is not None: + ld_overrides["quant_q_buffer"] = ( + forward_args.quant_q_buffer[ld_slice["gen_token_start"]: + ld_slice["gen_token_end"]]) + else: + # Slice the incoming offsets rather than rebuilding them from + # prompt lengths: for generation rows the prompt length is the + # history, not the number of packed Q rows. + for name in ("cu_q_seqlens", "cu_kv_seqlens"): + cu = getattr(forward_args, name) + if cu is not None: + ld_overrides[name] = _slice_locality_domain_cu_seqlens( + cu, ld_slice, forward_args.attention_input_type, + metadata.num_contexts) + if forward_args.fmha_scheduler_counter is not None: + ld_overrides["fmha_scheduler_counter"] = torch.zeros_like( + forward_args.fmha_scheduler_counter) + if forward_args.quant_q_buffer is not None: + ld_overrides["quant_q_buffer"] = cat( + forward_args.quant_q_buffer) + prepared.append( + dict(locality_domain_id=locality_domain_id, + ld_slice=ld_slice, + ld_args=replace(forward_args, **ld_overrides), + q=cat(q), + k=cat_kv(k), + v=cat_kv(v), + ld_output=ld_output, + ld_softmax_stats=ld_softmax_stats, + scatter_back=scatter_back, + num_ctx_tokens=num_ctx_tokens, + num_gen_tokens=num_gen_tokens)) + + num_kv_rows = k.shape[0] if k is not None else None + layout = (forward_args.attention_input_type, q.shape[0], num_kv_rows, + dense_ctx_kv, metadata.num_contexts, metadata.num_generations, + metadata.kv_lens_runtime.shape[0], + tuple(metadata.locality_domain_num_ctx_tokens), + tuple(metadata.locality_domain_ctx_req_splits), + tuple(metadata.locality_domain_gen_req_splits), + tuple(metadata.locality_domain_ctx_token_splits), + tuple(metadata.locality_domain_gen_token_splits)) + if metadata.locality_domain_validated_layout != layout: + _validate_locality_domain_coverage( + metadata, prepared, forward_args.attention_input_type, + q.shape[0], num_kv_rows, dense_ctx_kv) + metadata.locality_domain_validated_layout = layout + + two_stream = _LOCALITY_DOMAIN_TWO_STREAM + if two_stream: + runtime.fork() + try: + for prep in prepared: + locality_domain_id = prep["locality_domain_id"] + ld_slice = prep["ld_slice"] + # Serialized mode forks and joins around each domain in turn, so + # the domains never overlap. + if not two_stream: + runtime.fork() + try: + with _locality_domain_metadata(metadata, ld_slice, + locality_domain_id): + with runtime.partition_context(locality_domain_id): + yield (prep["q"], prep["k"], prep["v"], + prep["ld_args"]) + if prep["scatter_back"]: + _scatter_locality_domain_rows( + output, prep["ld_output"], ld_slice, + prep["num_ctx_tokens"], + prep["num_gen_tokens"]) + if prep["ld_softmax_stats"] is not None: + _scatter_locality_domain_rows( + forward_args.softmax_stats_tensor, + prep["ld_softmax_stats"], ld_slice, + prep["num_ctx_tokens"], + prep["num_gen_tokens"]) + if locality_domain_id == 0: + _resize_locality_domain_workspaces_from_domain0( + metadata) + finally: + if not two_stream: + runtime.join() + finally: + if two_stream: + runtime.join() + prepared.clear() + @classmethod def support_fused_rope(cls) -> bool: return True @@ -2984,6 +4085,28 @@ def _mla_rope_generation_impl( quant_scale_qkv, ) + # RoPE generation stays a single full-batch op; only the metadata the + # per-domain FMHA calls consume is split. + if (metadata.locality_domain_enabled + and metadata.kv_cache_manager is not None and getattr( + metadata.kv_cache_manager, "num_locality_domains", 1) > 1): + if (metadata.locality_domain_gen_req_splits is None + or metadata.locality_domain_gen_token_splits is None + or metadata.locality_domain_num_gen_tokens is None + or metadata.locality_domain_num_ctx_tokens is None): + raise RuntimeError( + "locality domain MLA generation needs the per-domain " + "splits that prepare() builds; they are missing") + _prepare_locality_domain_mla_gen_state_from_full( + metadata, + cu_q_seqlens, + cu_kv_seqlens, + fmha_scheduler_counter, + mla_bmm1_scale, + mla_bmm2_scale, + latent_cache.shape[0], + ) + def _fp4_mla_rope_generation( self, fused_q: torch.Tensor, diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/dummy_requests.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/dummy_requests.py index 5d6736365608..434fcfc9a8eb 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/dummy_requests.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/dummy_requests.py @@ -10,6 +10,7 @@ import torch +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState from tensorrt_llm._torch.pyexecutor.resource_manager import ( KVCacheManager, @@ -341,45 +342,80 @@ def create_cuda_graph_warmup_request( if context_requests is None: return None + num_long_requests = 1 + short_dummy_request_kwargs: dict[str, Any] = {} + long_dummy_request_kwargs: dict[str, Any] = {} + if ( + not num_contexts + and isinstance(kv_cache_manager, KVCacheManagerV2) + and kv_cache_manager.num_locality_domains > 1 + ): + num_locality_domains = kv_cache_manager.num_locality_domains + num_long_requests = min(num_locality_domains, batch_size) + # One longest request per domain grows every per-domain attention + # workspace to the long-sequence bound during capture. Short requests + # start from domain 0 to keep the generation-only split balanced. + short_dummy_request_kwargs["locality_domain_ids"] = [ + request_idx % num_locality_domains + for request_idx in range(batch_size - num_long_requests) + ] + long_dummy_request_kwargs["locality_domain_ids"] = list(range(num_long_requests)) + # Fill all rows but the longest with minimal-length generation requests. - num_fillers = batch_size - 1 - num_contexts - requests = kv_cache_manager.add_dummy_requests( - list(range(num_contexts, batch_size - 1)), - token_nums=( - None - if self._dummy_request_tokens is None - else [self._dummy_request_tokens] * num_fillers - ), - is_gen=True, - max_num_draft_tokens=runtime_draft_token_buffer_width, - kv_reserve_draft_tokens=self._config.max_draft_loop_tokens, - use_mrope=self._config.use_mrope, - max_beam_width=self._config.max_beam_width, - draft_kv_cache_manager=draft_kv_cache_manager, - capture_sampling_params=capture_sampling_params, - **self._dummy_request_kwargs(resource_manager, num_fillers), - ) - if requests is None: - for request in context_requests: - kv_cache_manager.free_resources(request) - if draft_kv_cache_manager is not None: - draft_kv_cache_manager.free_resources(request) - return None + num_fillers = batch_size - num_long_requests - num_contexts + requests = [] + if num_fillers > 0: + requests = kv_cache_manager.add_dummy_requests( + list(range(num_contexts, batch_size - num_long_requests)), + token_nums=( + None + if self._dummy_request_tokens is None + else [self._dummy_request_tokens] * num_fillers + ), + is_gen=True, + max_num_draft_tokens=runtime_draft_token_buffer_width, + kv_reserve_draft_tokens=self._config.max_draft_loop_tokens, + use_mrope=self._config.use_mrope, + max_beam_width=self._config.max_beam_width, + draft_kv_cache_manager=draft_kv_cache_manager, + capture_sampling_params=capture_sampling_params, + **self._dummy_request_kwargs(resource_manager, num_fillers), + **short_dummy_request_kwargs, + ) + if requests is None: + for request in context_requests: + kv_cache_manager.free_resources(request) + if draft_kv_cache_manager is not None: + draft_kv_cache_manager.free_resources(request) + return None requests = context_requests + requests - max_seq_len_request = self.add_longest_dummy_request( - resource_manager, - requests, - batch_size, - max_seq_len, - runtime_draft_token_buffer_width, - capture_sampling_params, - ) - if max_seq_len_request is None: + if num_long_requests == 1: + longest = self.add_longest_dummy_request( + resource_manager, + requests, + batch_size, + max_seq_len, + runtime_draft_token_buffer_width, + capture_sampling_params, + ) + max_seq_len_requests = None if longest is None else [longest] + else: + max_seq_len_requests = self.add_longest_dummy_requests( + resource_manager, + requests, + batch_size, + max_seq_len, + runtime_draft_token_buffer_width, + capture_sampling_params, + num_long_requests=num_long_requests, + dummy_request_kwargs=long_dummy_request_kwargs, + ) + if max_seq_len_requests is None: return None if num_contexts: - requests.append(max_seq_len_request) + requests.extend(max_seq_len_requests) for request, token_num in zip(requests, context_token_nums): request.state = LlmRequestState.CONTEXT_INIT request.context_current_position = 0 @@ -389,10 +425,9 @@ def create_cuda_graph_warmup_request( result.context_requests_last_chunk = requests[:num_contexts] result.generation_requests = requests[num_contexts:] else: - # Insert the longest request first to simulate padding for the CUDA - # graph. - requests.insert(0, max_seq_len_request) - result.generation_requests = requests + # Put the longest requests first to simulate padding for the CUDA + # graph; with locality domains this also seeds each domain group. + result.generation_requests = max_seq_len_requests + requests return self.finish_cuda_graph_warmup_request( result, resource_manager, batch_size, draft_len ) @@ -407,6 +442,28 @@ def add_longest_dummy_request( capture_sampling_params: SamplingParams | None, ) -> LlmRequest | None: """Add the dummy request with the longest sequence; free ``requests`` if it does not fit.""" + longest = self.add_longest_dummy_requests( + resource_manager, + requests, + batch_size, + max_seq_len, + runtime_draft_token_buffer_width, + capture_sampling_params, + ) + return None if longest is None else longest[0] + + def add_longest_dummy_requests( + self, + resource_manager: ResourceManager, + requests: list[LlmRequest], + batch_size: int, + max_seq_len: int | None, + runtime_draft_token_buffer_width: int, + capture_sampling_params: SamplingParams | None, + num_long_requests: int = 1, + dummy_request_kwargs: dict[str, Any] | None = None, + ) -> list[LlmRequest] | None: + """Add the longest dummy requests; free ``requests`` if they do not fit.""" kv_cache_manager = resource_manager.get_resource_manager(self.kv_cache_manager_key) draft_kv_cache_manager = get_draft_kv_cache_manager( self._config.spec_config, resource_manager @@ -418,7 +475,8 @@ def free_warmup_requests() -> None: if draft_kv_cache_manager is not None: draft_kv_cache_manager.free_resources(r) - # Add one dummy request with the maximum possible sequence length. + # Add the dummy requests with the maximum possible sequence length. + available_tokens_batch_size = math.ceil(batch_size / num_long_requests) max_seq_len = min( self._config.max_seq_len if max_seq_len is None else max_seq_len, kv_cache_manager.max_seq_len, @@ -428,13 +486,15 @@ def free_warmup_requests() -> None: # for the actual KV reservation per request. _kv_draft = self._config.max_draft_loop_tokens available_tokens = kv_cache_manager.get_num_available_tokens( - token_num_upper_bound=max_seq_len, batch_size=batch_size, max_num_draft_tokens=_kv_draft + token_num_upper_bound=max_seq_len, + batch_size=available_tokens_batch_size, + max_num_draft_tokens=_kv_draft, ) # Also consider draft KV cache capacity when it exists if draft_kv_cache_manager is not None: draft_available_tokens = draft_kv_cache_manager.get_num_available_tokens( - batch_size=batch_size, + batch_size=available_tokens_batch_size, token_num_upper_bound=max_seq_len, max_num_draft_tokens=_kv_draft, ) @@ -453,9 +513,9 @@ def free_warmup_requests() -> None: token_num = int(token_num) # Ensure int for range() in add_dummy_requests - max_seq_len_request = kv_cache_manager.add_dummy_requests( - request_ids=[batch_size - 1], - token_nums=[token_num], + max_seq_len_requests = kv_cache_manager.add_dummy_requests( + request_ids=list(range(batch_size - num_long_requests, batch_size)), + token_nums=[token_num] * num_long_requests, is_gen=True, max_num_draft_tokens=runtime_draft_token_buffer_width, kv_reserve_draft_tokens=self._config.max_draft_loop_tokens, @@ -463,13 +523,14 @@ def free_warmup_requests() -> None: max_beam_width=self._config.max_beam_width, draft_kv_cache_manager=draft_kv_cache_manager, capture_sampling_params=capture_sampling_params, - **self._dummy_request_kwargs(resource_manager, 1), + **self._dummy_request_kwargs(resource_manager, num_long_requests), + **(dummy_request_kwargs or {}), ) - if max_seq_len_request is None: + if max_seq_len_requests is None: free_warmup_requests() return None - return max_seq_len_request[0] + return max_seq_len_requests def finish_cuda_graph_warmup_request( self, diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py index 6c409f0eaf54..fffc43636fa5 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py @@ -58,6 +58,7 @@ from ..common import get_top_level_model, prepare_multimodal_indices from .config import DecoderRunnerConfig from .dummy_requests import DummyRequestFactory +from .locality_domain import locality_domain_forkjoin_enabled from .metadata import DecoderMetadata from .warmup import WarmupAgreement @@ -456,6 +457,9 @@ def promote_final_contexts( # fresh final-context row would skip its cache-slot reset. and not self._model_uses_ple_recurrent_state and self.mapping.cp_size == 1 + # Promoted rows are re-sorted by domain with the generation rows, + # while sampling still walks the semantic batch in its own order. + and not locality_domain_forkjoin_enabled(kv_cache_manager) ): return scheduled_requests, frozenset() return _make_single_token_context_graph_batch( diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py index 578ab3e68667..dc6733d6e01a 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py @@ -7,6 +7,7 @@ import os from collections.abc import Callable, Sequence from dataclasses import dataclass, field, fields, replace +from itertools import accumulate from typing import Any import torch @@ -62,6 +63,11 @@ ) from ..interface import PreparedInputs from .config import DecoderRunnerConfig +from .locality_domain import ( + locality_domain_forkjoin_enabled, + locality_domain_row_splits, + sort_generation_requests_by_locality_domain, +) from .metadata import set_spec_metadata_all_rank_num_tokens @@ -274,6 +280,8 @@ class HostInputs: mrope_dummy_seq_slot: int extra_inputs: ExtraInputsCollector context_prompt_lookahead: list[int] | None + # Fork-join attention needs each domain's rows contiguous in the packed batch. + locality_domain_reorder: bool = False lists: HostLists = field(default_factory=HostLists) num_ctx_requests: int = 0 @@ -810,6 +818,7 @@ def prepare_inputs( else None ), ) + host.locality_domain_reorder = locality_domain_forkjoin_enabled(kv_cache_manager) self._pack_context_requests(host, scheduled_requests) self._classify_generation_requests( host, @@ -979,6 +988,16 @@ def _pack_context_requests( extra_inputs = host.extra_inputs context_prompt_lookahead = host.context_prompt_lookahead + # Chunked prefill keeps the scheduler order: MLA context then runs on + # the single-kernel path and the context splits fall back to positional. + if ( + host.locality_domain_reorder + and not self._config.attention_runtime_features.chunked_prefill + ): + scheduled_requests.reset_context_requests( + sorted(scheduled_requests.context_requests, key=lambda r: r.py_locality_domain_id) + ) + for request in scheduled_requests.context_requests: request_ids.append(request.py_request_id) draft_lens.append(0) @@ -1180,6 +1199,13 @@ def _classify_generation_requests( generation_requests = host.lists.generation_requests all_gen_request_ids = host.lists.all_gen_request_ids + # Reordering in place keeps the sampler and spec-decode on the same + # order as the model input, so no logit un-permutation is needed. + if host.locality_domain_reorder and scheduled_requests.generation_requests: + scheduled_requests.generation_requests = sort_generation_requests_by_locality_domain( + list(scheduled_requests.generation_requests), promoted_context_request_ids + ) + # Requests with draft tokens are treated like extend requests. Dummy extend requests should be # at the end of extend_requests. for request in scheduled_requests.generation_requests: @@ -2094,6 +2120,10 @@ def _prepare_attn_metadata( ), ) attn_metadata.kv_cache_manager = kv_cache_manager + if host.locality_domain_reorder: + self._set_locality_domain_splits( + host, scheduled_requests, kv_cache_manager, attn_metadata + ) if hasattr(self.model.model_config.pretrained_config, "chunk_size"): attn_metadata.mamba_chunk_size = self.model.model_config.pretrained_config.chunk_size @@ -2115,6 +2145,56 @@ def _prepare_attn_metadata( host.num_cached_tokens_snapshot = num_cached_tokens_snapshot host.remote_tail_modes = remote_tail_modes + def _set_locality_domain_splits( + self, + host: HostInputs, + scheduled_requests: ScheduledRequests, + kv_cache_manager: KVCacheManagerV2, + attn_metadata: AttentionMetadata, + ) -> None: + """Record the per-domain request and token boundaries for fork-join attention. + + The single ``attn_metadata.kv_cache_manager`` is shared across domains: + V2 addresses every domain through global block indices. + """ + attn_metadata.locality_domain_enabled = True + num_locality_domains = kv_cache_manager.num_locality_domains + sequence_lengths = host.lists.sequence_lengths + num_ctx_requests = host.num_ctx_requests + + # Context rows keep the scheduler order under chunked prefill, so the + # split follows ownership only when the domains are contiguous. + per_locality_domain_ctx_reqs, per_locality_domain_ctx_tokens = locality_domain_row_splits( + scheduled_requests.context_requests, + sequence_lengths[:num_ctx_requests], + num_locality_domains, + positional=False, + ) + + # Rows are packed extend then generation, which is not the domain-sorted + # order once the categories mix, so count from the packed list. + per_locality_domain_gen_reqs, per_locality_domain_gen_tokens = locality_domain_row_splits( + host.lists.extend_requests + host.lists.generation_requests, + sequence_lengths[num_ctx_requests:], + num_locality_domains, + positional=num_ctx_requests == 0, + ) + + attn_metadata.locality_domain_ctx_req_splits = [0] + list( + accumulate(per_locality_domain_ctx_reqs) + ) + attn_metadata.locality_domain_gen_req_splits = [0] + list( + accumulate(per_locality_domain_gen_reqs) + ) + attn_metadata.locality_domain_ctx_token_splits = [0] + list( + accumulate(per_locality_domain_ctx_tokens) + ) + attn_metadata.locality_domain_gen_token_splits = [0] + list( + accumulate(per_locality_domain_gen_tokens) + ) + attn_metadata.locality_domain_num_ctx_tokens = per_locality_domain_ctx_tokens + attn_metadata.locality_domain_num_gen_tokens = per_locality_domain_gen_tokens + def _pad_inputs( self, host: HostInputs, diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/locality_domain.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/locality_domain.py new file mode 100644 index 000000000000..6b49e2968204 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/locality_domain.py @@ -0,0 +1,99 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Locality domain request ordering and split arithmetic for the decoder runner.""" + +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest + + +def locality_domain_forkjoin_enabled(kv_cache_manager) -> bool: + """Whether to drive two-stream attention for *kv_cache_manager*. + + True iff the manager is :class:`KVCacheManagerV2` with + ``num_locality_domains > 1`` and ``fork_join_attn`` set. + """ + if not isinstance(kv_cache_manager, KVCacheManagerV2): + return False + if kv_cache_manager.num_locality_domains <= 1: + return False + return kv_cache_manager.fork_join_attn + + +def sort_generation_requests_by_locality_domain( + requests: list[LlmRequest], promoted_context_request_ids: frozenset[int] = frozenset() +) -> list[LlmRequest]: + """Group by locality domain inside the overlap scheduler's packing order. + + Input tokens are packed as [rows whose token comes from the host][rows that + read it from the previous step's output], while positions follow request + order, so host-token rows must stay first. CUDA graph padding dummies stay + at the tail so ``pad_batch``'s ``generation_requests[:-padding_size]`` trims + only them. + """ + + def reads_previous_step(request: LlmRequest) -> bool: + return not ( + request.py_request_id in promoted_context_request_ids + or request.is_dummy + or request.py_batch_idx is None + ) + + return sorted( + requests, + key=lambda r: (r.is_cuda_graph_dummy, reads_previous_step(r), r.py_locality_domain_id), + ) + + +def balanced_generation_dispatch(num_rows: int, num_locality_domains: int) -> list[int]: + """Domain for each packed generation row, by position only. + + CUDA graph capture bakes the per-domain request counts in as launch scalars, + so the split must depend on the padded row count alone, never on which + domain owns a row: capturing (half-2, half+2) and replaying (half+1, half-1) + crashes. Rows that are repeated padding dummies count once per occurrence. + """ + base, extra = divmod(num_rows, num_locality_domains) + dispatch: list[int] = [] + for domain in range(num_locality_domains): + dispatch.extend([domain] * (base + (1 if domain < extra else 0))) + return dispatch + + +def locality_domain_row_splits( + packed_requests: list[LlmRequest], + sequence_lengths: list[int], + num_locality_domains: int, + *, + positional: bool, +) -> tuple[list[int], list[int]]: + """Per-domain request and token counts over one packed row range. + + The attention side turns the counts into cumulative ranges, so ownership + can only be followed when the packed order is sorted by domain. Mixed + generation categories (extend, generation) and unsorted context rows under + chunked prefill break that, and then the rows are split by position + instead, which is still correct and only costs locality. + """ + if not positional: + dispatch = [request.py_locality_domain_id for request in packed_requests] + positional = dispatch != sorted(dispatch) + if positional: + dispatch = balanced_generation_dispatch(len(packed_requests), num_locality_domains) + per_domain_reqs = [0] * num_locality_domains + per_domain_tokens = [0] * num_locality_domains + for i, domain in enumerate(dispatch): + per_domain_reqs[domain] += 1 + per_domain_tokens[domain] += sequence_lengths[i] + return per_domain_reqs, per_domain_tokens diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index bac560ef8f36..e7f01f2ee8d8 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -794,8 +794,10 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, # TrtllmAttentionMetadata, and sparse metadata all derive from it. draft_scope.callback(_reset_mtp_index_share, attn_metadata) draft_attn_metadata = draft_scope.enter_context( - self.draft_kv_cache_context(attn_metadata, - draft_kv_cache_manager)) + self.draft_kv_cache_context( + attn_metadata, + draft_kv_cache_manager, + enable_draft_locality_domain_forkjoin=True)) attn_metadata = draft_attn_metadata inputs["attn_metadata"] = draft_attn_metadata if self.is_mtp_eagle and is_dsa_cache_manager( diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index eb9067c54f95..42f2a132cc88 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -3184,19 +3184,40 @@ def get_draft_kv_cache_manager(self, resource_manager): return None @contextmanager - def draft_kv_cache_context(self, attn_metadata, draft_kv_cache_manager): + def draft_kv_cache_context( + self, + attn_metadata, + draft_kv_cache_manager, + enable_draft_locality_domain_forkjoin: bool = False): """ - Select draft attention metadata for one-engine speculative decoding. + Context manager for the draft sub-forward in one-engine speculative decoding. TRTLLM metadata temporarily swaps its manager and cache-layout-dependent buffers, including DSA indexer offsets and slot mappings. FlashInfer uses an independently planned metadata view because its page tables and kernel wrappers are manager-specific. + + When the caller opts in, keep fork-join/two-stream attention enabled. + For separate-draft topology, also swap + ``locality_domain_kv_cache_block_offsets`` to the draft cache's + per-locality domain offsets so kernels address the draft KV layout. """ - # draft_kv_cache_manager is None if using two-engine speculative decoding or not enabling separate draft KV cache. + # draft_kv_cache_manager is None if using two-engine speculative decoding + # or not enabling separate draft KV cache. if draft_kv_cache_manager is None: - yield attn_metadata + if (enable_draft_locality_domain_forkjoin or not getattr( + attn_metadata, 'locality_domain_enabled', False)): + yield attn_metadata + return + # Draft workers that rewrite seq_lens/num_contexts in place (dynamic + # trees) leave the per-domain splits stale; without an explicit + # opt-in the draft sub-forward runs unpartitioned. + attn_metadata.locality_domain_enabled = False + try: + yield attn_metadata + finally: + attn_metadata.locality_domain_enabled = True return from ..attention.backends.flashinfer import FlashInferAttentionMetadata @@ -3210,13 +3231,39 @@ def draft_kv_cache_context(self, attn_metadata, draft_kv_cache_manager): saved_state = prepare_attn_metadata_for_draft_replay( attn_metadata, draft_kv_cache_manager) - if saved_state is None: - yield attn_metadata - return + saved_locality_attrs = [] + + def patch_locality_metadata(name, value): + saved_locality_attrs.append((name, getattr(attn_metadata, name))) + setattr(attn_metadata, name, value) + draft_manager_supports_locality_domain_forkjoin = getattr( + draft_kv_cache_manager, 'fork_join_attn', False) + using_separate_draft_cache = saved_state is not None try: + if attn_metadata.locality_domain_enabled: + if not enable_draft_locality_domain_forkjoin: + patch_locality_metadata('locality_domain_enabled', False) + elif using_separate_draft_cache: + if draft_manager_supports_locality_domain_forkjoin: + draft_locality_domain_block_offsets = ( + attn_metadata. + draft_locality_domain_kv_cache_block_offsets) + if draft_locality_domain_block_offsets is None: + raise RuntimeError( + "locality domain fork-join is enabled for the draft KV " + "cache, but draft_locality_domain_kv_cache_block_offsets " + "was not prepared.") + patch_locality_metadata( + 'locality_domain_kv_cache_block_offsets', + draft_locality_domain_block_offsets) + else: + patch_locality_metadata('locality_domain_enabled', + False) yield attn_metadata finally: + for name, value in reversed(saved_locality_attrs): + setattr(attn_metadata, name, value) restore_attn_metadata_after_draft_replay(attn_metadata, saved_state) def _sample_tokens_for_batch( diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index 2107fbed26c8..997a1c2b91f9 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -472,7 +472,10 @@ def _forward_impl( draft_kv_cache_manager = self.get_draft_kv_cache_manager( resource_manager) - with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): + with self.draft_kv_cache_context( + attn_metadata, + draft_kv_cache_manager, + enable_draft_locality_domain_forkjoin=True): for i, mtp_layer in enumerate( draft_model.mtp_layers[:runtime_draft_len]): if self.guided_decoder is not None: diff --git a/tests/integration/defs/kv_cache/test_stream_localization_e2e.py b/tests/integration/defs/kv_cache/test_stream_localization_e2e.py new file mode 100644 index 000000000000..8587bd10193f --- /dev/null +++ b/tests/integration/defs/kv_cache/test_stream_localization_e2e.py @@ -0,0 +1,99 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Integration tests for stream localization equivalence. + +Verifies that running inference with locality domains enabled produces +identical outputs to running without them, using the LLM API end-to-end. + +Rubin only: locality domains need hardware that exposes two locality domains. +``TRT_LLM_MOCK_LOCALIZATION_SUPPORT`` is deliberately not used here -- it only +reaches the KV cache allocator, so the fork/join runtime would still fail. + +Model: Qwen3-4B. +""" + +from ..conftest import llm_models_root, skip_no_rubin + +# Qwen3-4B: dense GQA, so attention arrives as a single mixed call, a path +# an MLA model never exercises. +_MODEL_PATH = f"{llm_models_root()}/Qwen3/Qwen3-4B" + +SHORT_PROMPTS = [ + "What is 2+2? Answer in one number.", + "Capital of France? One word.", + "Largest planet in our solar system? One word.", + "Who wrote Romeo and Juliet? One name.", +] + +MIXED_LENGTH_PROMPTS = [ + "Hi", + "The quick brown fox jumps over the lazy dog and then runs across the field", + "Once upon a time in a land far far away there lived a brave knight who fought dragons every day", + "A", +] + + +def _generate(model_path, prompts, *, localized: bool): + """Run generation with or without locality domains.""" + from tensorrt_llm import LLM + from tensorrt_llm.llmapi import KvCacheConfig, SamplingParams, SchedulerConfig + + kv_config = KvCacheConfig(use_kv_cache_manager_v2=True) + scheduler_config = SchedulerConfig(capacity_scheduler_policy="MAX_UTILIZATION") + sampling = SamplingParams(temperature=0.0, max_tokens=32) + + with LLM( + model_path, + kv_cache_config=kv_config, + scheduler_config=scheduler_config, + # Without this the localized arm runs ordinary mode and the comparison + # holds no matter what locality domains do. The capability is probed + # from the device rather than mocked: the mock only reaches the KV cache + # allocator, so the fork/join runtime still requires real hardware. + enable_locality_domains=localized, + ) as llm: + return llm.generate(prompts, sampling_params=sampling) + + +def _assert_outputs_match(baseline, localized, prompts): + """Assert that baseline and localized outputs produce identical text.""" + assert len(baseline) == len(localized) == len(prompts) + for i, (base_output, localized_output) in enumerate(zip(baseline, localized)): + assert base_output.outputs[0].text == localized_output.outputs[0].text, ( + f"Prompt {i}: baseline={base_output.outputs[0].text!r} " + f"vs localized={localized_output.outputs[0].text!r}" + ) + + +@skip_no_rubin +class TestStreamLocalizationEquivalence: + """Verify that localized (split->forward->merge) matches single-batch forward. + + Uses the LLM API with Qwen3-4B to run the same prompts with and + without locality domains, comparing generated text at temperature=0. + """ + + def test_generation_only_equivalence(self): + """Short prompts -- all requests complete context quickly.""" + baseline = _generate(_MODEL_PATH, SHORT_PROMPTS, localized=False) + localized = _generate(_MODEL_PATH, SHORT_PROMPTS, localized=True) + _assert_outputs_match(baseline, localized, SHORT_PROMPTS) + + def test_mixed_length_equivalence(self): + """Varying prompt lengths -- exercises context chunking + generation.""" + baseline = _generate(_MODEL_PATH, MIXED_LENGTH_PROMPTS, localized=False) + localized = _generate(_MODEL_PATH, MIXED_LENGTH_PROMPTS, localized=True) + _assert_outputs_match(baseline, localized, MIXED_LENGTH_PROMPTS) diff --git a/tests/unittest/_torch/attention/test_attention_mla.py b/tests/unittest/_torch/attention/test_attention_mla.py index 83133dc0c924..1236ef0844c2 100644 --- a/tests/unittest/_torch/attention/test_attention_mla.py +++ b/tests/unittest/_torch/attention/test_attention_mla.py @@ -1394,6 +1394,11 @@ def __init__(self, q_lens, kv_lens, num_contexts=0): device=device) # `enable_flash_mla` is the only other attribute the two hooks read. self.enable_flash_mla = False + # _invalidate_mla_scheduler_buffers consults these before recomputing + # the per-locality-domain splits; False short-circuits that branch. + self.locality_domain_enabled = False + self.kv_cache_manager = None + self.request_ids = None self._invalidate_mla_scheduler_buffers() def rebuild(self): diff --git a/tests/unittest/_torch/attention/test_locality_domain_attention_gpu.py b/tests/unittest/_torch/attention/test_locality_domain_attention_gpu.py new file mode 100644 index 000000000000..d1052a5f2f46 --- /dev/null +++ b/tests/unittest/_torch/attention/test_locality_domain_attention_gpu.py @@ -0,0 +1,306 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Decode attention with locality domains on matches it with them off (Rubin only). + +The same requests, cached K/V and query run through ``TrtllmAttention`` twice: once on an ordinary +KV cache with one kernel over the batch, once on a localized KV cache where each locality domain's +rows run on that domain's SMs. Per-domain kernels choose their configuration for the domain's SMs, +so the two runs may use different kernel variants and KV splits; the outputs must still agree. +""" + +import math +from types import SimpleNamespace + +import pytest +import torch +from kv_cache_utils import fill_kv_cache_logical + +import tensorrt_llm +from tensorrt_llm._torch.attention.backends import TrtllmAttention +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionInputType, + MLAParams, + PositionalEmbeddingParams, + PredefinedAttentionMask, + RopeParams, +) +from tensorrt_llm._torch.locality_domain_utils import is_locality_domain_enabled +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm.bindings import DataType +from tensorrt_llm.functional import PositionEmbeddingType +from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or not is_locality_domain_enabled(), + reason="needs a GPU with locality domains (Rubin)", +) + +_KV_LORA_RANK = 512 +_ROPE_DIM = 64 +_MLA_HEAD_DIM = _KV_LORA_RANK + _ROPE_DIM + + +def _domains(batch: int, localized: bool): + # Domain 0 first, so the positional split runs every row on the domain holding its KV. + if not localized: + return None + half = (batch + 1) // 2 + return [0] * half + [1] * (batch - half) + + +def _manager( + *, mla: bool, batch: int, cached: int, num_kv_heads: int, head_dim: int, localized: bool +): + tokens_per_block = 32 if mla else 64 + max_seq_len = cached + 1 + tokens_per_block + blocks_per_seq = math.ceil(max_seq_len / tokens_per_block) + cache_type = tensorrt_llm.bindings.internal.batch_manager.CacheType + manager = KVCacheManagerV2( + # Each domain holds half the pool, so leave room for a full batch in either one. + KvCacheConfig( + max_tokens=blocks_per_seq * tokens_per_block * 2 * (batch + 2), enable_block_reuse=False + ), + cache_type.SELFKONLY if mla else cache_type.SELF, + num_layers=1, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + max_batch_size=batch, + mapping=Mapping(world_size=1, tp_size=1, rank=0), + dtype=DataType.FP8 if mla else DataType.BF16, + enable_locality_domains=localized, + ) + requests = manager.add_dummy_requests( + list(range(batch)), [cached + 1] * batch, locality_domain_ids=_domains(batch, localized) + ) + assert requests is not None + if localized: + assert manager.num_locality_domains == 2 + return manager + + +def _metadata(manager, *, mla: bool, batch: int, cached: int, localized: bool): + kwargs = dict( + num_contexts=0, + kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[cached] * batch), + seq_lens=torch.ones(batch, dtype=torch.int), + max_num_requests=batch, + max_num_tokens=8192, + kv_cache_manager=manager, + request_ids=list(range(batch)), + prompt_lens=[cached + 1] * batch, + mapping=Mapping(world_size=1, tp_size=1, rank=0), + locality_domain_enabled=localized, + ) + if mla: + kwargs["enable_flash_mla"] = False + metadata = TrtllmAttention.Metadata(**kwargs) + metadata.prepare() + return metadata + + +def _fill_mla_latent(manager, latent_per_seq): + """Write each request's cached latent rows through its own page list (layout-independent).""" + buf = manager.get_buffers(0, kv_layout="NHD") # [pages, 1, tokens_per_block, 1, 576] + tokens_per_block = buf.shape[2] + for blocks, latent in zip( + manager.get_batch_cache_indices(list(range(len(latent_per_seq))), 0), latent_per_seq + ): + blocks = [b for b in blocks if b != -1] + # Zero only owned pages: a localized pool leaves unmapped VA between the domains' pages. + buf[blocks] = 0 + written = 0 + for blk in blocks: + if written >= latent.shape[0]: + break + n = min(tokens_per_block, latent.shape[0] - written) + buf[blk, 0, :n, 0].copy_(latent[written : written + n].to(buf.dtype)) + written += n + + +def _gqa_output(*, batch: int, cached: int, localized: bool) -> torch.Tensor: + num_heads, num_kv_heads, head_dim = 48, 4, 128 + gen = torch.Generator(device="cuda").manual_seed(0) + k = [ + torch.randn(cached, num_kv_heads, head_dim, device="cuda", generator=gen) * 0.5 + for _ in range(batch) + ] + v = [ + torch.randn(cached, num_kv_heads, head_dim, device="cuda", generator=gen) + for _ in range(batch) + ] + qkv = torch.randn( + batch, (num_heads + 2 * num_kv_heads) * head_dim, device="cuda", generator=gen + ).to(torch.bfloat16) + manager = _manager( + mla=False, + batch=batch, + cached=cached, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + localized=localized, + ) + try: + fill_kv_cache_logical(manager, 0, list(range(batch)), k, v, kv_layout="HND") + metadata = _metadata(manager, mla=False, batch=batch, cached=cached, localized=localized) + backend = TrtllmAttention( + layer_idx=0, num_heads=num_heads, num_kv_heads=num_kv_heads, head_dim=head_dim + ) + mask = PredefinedAttentionMask.CAUSAL + output = backend.create_output( + qkv, is_quantize_output=False, metadata=metadata, attention_mask=mask, is_gen_only=True + )[0] + backend.forward( + qkv, + None, + None, + metadata, + attention_mask=mask, + attention_input_type=AttentionInputType.generation_only, + output=output, + ) + torch.cuda.synchronize() + return output.float().cpu() + finally: + manager.shutdown() + + +def _mla_output(*, batch: int, cached: int, localized: bool) -> torch.Tensor: + num_heads = 64 + gen = torch.Generator(device="cuda").manual_seed(0) + latent = [ + torch.randn(cached, _MLA_HEAD_DIM, device="cuda", generator=gen) * 0.5 for _ in range(batch) + ] + fused_q = torch.randn(batch, num_heads * _MLA_HEAD_DIM, device="cuda", generator=gen).to( + torch.bfloat16 + ) + q_pe = torch.randn(batch, num_heads, _ROPE_DIM, device="cuda", generator=gen).to(torch.bfloat16) + new_latent = torch.randn(batch, _MLA_HEAD_DIM, device="cuda", generator=gen).to(torch.bfloat16) + manager = _manager( + mla=True, + batch=batch, + cached=cached, + num_kv_heads=1, + head_dim=_MLA_HEAD_DIM, + localized=localized, + ) + try: + _fill_mla_latent(manager, latent) + metadata = _metadata(manager, mla=True, batch=batch, cached=cached, localized=localized) + # DeepSeek-V3 yarn RoPE, as in test_attention_mla.py. + rope = RopeParams.from_config( + SimpleNamespace( + hidden_size=7168, + num_attention_heads=128, + rope_scaling={ + "beta_fast": 32, + "beta_slow": 1, + "factor": 40.0, + "mscale": 1.0, + "mscale_all_dim": 1.0, + "original_max_position_embeddings": 4096, + "type": "yarn", + }, + max_position_embeddings=163840, + rope_theta=10000.0, + qk_rope_head_dim=_ROPE_DIM, + model_type="deepseek_v3", + ) + ) + backend = TrtllmAttention( + layer_idx=0, + num_heads=num_heads, + head_dim=_MLA_HEAD_DIM, + num_kv_heads=1, + quant_config=QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8.value), + pos_embd_params=PositionalEmbeddingParams( + type=PositionEmbeddingType.yarn, rope=rope, is_neox=False + ), + mla_params=MLAParams( + q_lora_rank=1536, + kv_lora_rank=_KV_LORA_RANK, + qk_rope_head_dim=_ROPE_DIM, + qk_nope_head_dim=128, + v_head_dim=128, + predicted_tokens_per_seq=1, + ), + ) + cu_q = torch.empty(batch + 1, dtype=torch.int32, device="cuda") + cu_kv = torch.empty(batch + 1, dtype=torch.int32, device="cuda") + counter = torch.empty(1, dtype=torch.uint32, device="cuda") + bmm1_scale = torch.empty(2, dtype=torch.float32, device="cuda") + bmm2_scale = torch.empty(1, dtype=torch.float32, device="cuda") + quant_q = torch.empty(batch, num_heads * _MLA_HEAD_DIM, dtype=torch.uint8, device="cuda") + backend.mla_rope_generation( + fused_q, + q_pe, + new_latent, + metadata, + cu_q, + cu_kv, + counter, + bmm1_scale, + bmm2_scale, + quant_q, + ) + output = backend.create_output( + fused_q, + is_quantize_output=False, + metadata=metadata, + attention_mask=PredefinedAttentionMask.CAUSAL, + is_gen_only=True, + )[0] + backend.forward( + fused_q, + None, + None, + metadata, + attention_input_type=AttentionInputType.generation_only, + latent_cache=new_latent, + q_pe=q_pe, + cu_q_seqlens=cu_q, + cu_kv_seqlens=cu_kv, + fmha_scheduler_counter=counter, + mla_bmm1_scale=bmm1_scale, + mla_bmm2_scale=bmm2_scale, + quant_q_buffer=quant_q, + output=output, + ) + torch.cuda.synchronize() + return output.float().cpu() + finally: + manager.shutdown() + + +@pytest.mark.parametrize("batch,cached", [(8, 1000), (64, 4096)]) +def test_gqa_decode_matches_with_locality_domains(batch, cached): + off = _gqa_output(batch=batch, cached=cached, localized=False) + on = _gqa_output(batch=batch, cached=cached, localized=True) + assert torch.isfinite(on).all() + torch.testing.assert_close(on, off, atol=2e-2, rtol=2e-2) + + +# At 64 x 16k each domain's kernel sizes its KV split for half the SMs, unlike the single full-device kernel. +@pytest.mark.parametrize("batch,cached", [(8, 1000), (64, 16384)]) +def test_mla_decode_matches_with_locality_domains(batch, cached): + off = _mla_output(batch=batch, cached=cached, localized=False) + on = _mla_output(batch=batch, cached=cached, localized=True) + assert torch.isfinite(on).all() + torch.testing.assert_close(on, off, atol=2e-2, rtol=2e-2) diff --git a/tests/unittest/_torch/attention/test_locality_domain_partitions.py b/tests/unittest/_torch/attention/test_locality_domain_partitions.py new file mode 100644 index 000000000000..5c5761772d42 --- /dev/null +++ b/tests/unittest/_torch/attention/test_locality_domain_partitions.py @@ -0,0 +1,435 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""CPU checks for the locality-domain partition helpers of the TRTLLM backend.""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends import trtllm as trtllm_backend +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + AttentionInputType, +) +from tensorrt_llm._torch.speculative.interface import SpecWorkerBase + +pytestmark = pytest.mark.cpu_only + + +class _Metadata: + """Stand-in exposing the fields the partition helpers touch.""" + + def __init__(self, num_ctx, num_gen, ctx_tokens=(), gen_tokens=(), use_spec_decoding=False): + self._num_contexts = num_ctx + self._num_generations = num_gen + self._num_ctx_tokens = sum(ctx_tokens) + self._num_tokens = sum(ctx_tokens) + sum(gen_tokens) + n = num_ctx + num_gen + self.kv_lens_cuda_runtime = torch.arange(n) + self.kv_lens_runtime = torch.arange(n) + self.prompt_lens_cuda_runtime = torch.arange(n) + self.prompt_lens_cpu_runtime = torch.arange(n) + self.host_request_types_runtime = torch.zeros(n, dtype=torch.int) + self.host_total_kv_lens = torch.zeros(2, dtype=torch.int) + self.kv_cache_block_offsets = torch.zeros(1, n, 2, 4, dtype=torch.int) + self.block_ids_per_seq = torch.zeros(n, 4, dtype=torch.int) + self.locality_domain_workspaces = None + self.locality_domain_cuda_graph_workspaces = None + self.locality_domain_spec_decoding_bl_tree_masks = None + self.locality_domain_num_ctx_tokens = list(ctx_tokens) + self.use_spec_decoding = use_spec_decoding + self.spec_decoding_generation_lengths = None + self.spec_decoding_packed_mask = None + self.spec_decoding_bl_tree_mask_offset = None + self.spec_bl_tree_first_sparse_mask_offset_kv = None + self.spec_decoding_bl_tree_mask = None + self.spec_decoding_position_offsets = None + self.spec_decoding_position_offsets_cpp = None + self.spec_decoding_position_offsets_for_cpp = None + + @property + def num_contexts(self): + return self._num_contexts + + @property + def num_generations(self): + return self._num_generations + + +def _slice( + *, + ctx_req=(0, 0), + gen_req=(0, 0), + ctx_tok=(0, 0), + gen_tok=(0, 0), + ctx_kv=None, + n=1, +): + ctx_kv = ctx_tok if ctx_kv is None else ctx_kv + num_ctx_tokens = ctx_tok[1] - ctx_tok[0] + num_gen_tokens = gen_tok[1] - gen_tok[0] + return { + "ctx_req_start": ctx_req[0], + "ctx_req_end": ctx_req[1], + "gen_req_start": gen_req[0], + "gen_req_end": gen_req[1], + "ctx_token_start": ctx_tok[0], + "ctx_token_end": ctx_tok[1], + "ctx_kv_token_start": ctx_kv[0], + "ctx_kv_token_end": ctx_kv[1], + "gen_token_start": gen_tok[0], + "gen_token_end": gen_tok[1], + "num_ctx_reqs": ctx_req[1] - ctx_req[0], + "num_gen_reqs": gen_req[1] - gen_req[0], + "num_ctx_tokens": num_ctx_tokens, + "num_gen_tokens": num_gen_tokens, + "num_tokens": num_ctx_tokens + num_gen_tokens, + "num_seqs": (ctx_req[1] - ctx_req[0]) + (gen_req[1] - gen_req[0]), + "sequence_lengths": torch.ones(n), + "kv_lengths": torch.ones(n), + "prompt_lengths_cuda": torch.ones(n), + "prompt_lengths_cpu": torch.ones(n), + "host_context_lengths": torch.ones(n), + "request_types": torch.ones(n, dtype=torch.int), + "host_total_kv_lengths": torch.ones(2, dtype=torch.int), + "block_offsets": torch.ones(1, n, 2, 4, dtype=torch.int), + } + + +def _prepared(*slices): + return [{"ld_slice": s} for s in slices] + + +def _gate_metadata(enabled=True, num_domains=2, **extra): + fields = dict( + locality_domain_enabled=enabled, + kv_cache_manager=SimpleNamespace(num_locality_domains=num_domains), + locality_domain_gen_req_splits=[0, 1, 2], + runtime_features=None, + beam_width=1, + is_cross=False, + enable_helix=False, + ) + fields.update(extra) + return SimpleNamespace(**fields) + + +_ATTN = SimpleNamespace(is_mla_enable=False) + + +def test_gate_requires_metadata_opt_in(): + forward_args = AttentionForwardArgs(output=torch.empty(1)) + assert not trtllm_backend._locality_domain_attention_enabled( + _ATTN, _gate_metadata(enabled=False), forward_args + ) + assert trtllm_backend._locality_domain_attention_enabled(_ATTN, _gate_metadata(), forward_args) + assert not trtllm_backend._locality_domain_attention_enabled( + _ATTN, _gate_metadata(num_domains=1), forward_args + ) + assert not trtllm_backend._locality_domain_attention_enabled( + _ATTN, _gate_metadata(locality_domain_gen_req_splits=None), forward_args + ) + + +@pytest.mark.parametrize( + "metadata_extra, args_extra", + [ + ({"beam_width": 2}, {}), + ({"is_cross": True}, {}), + ({"enable_helix": True}, {}), + ({}, {"attention_mask_data": torch.zeros(1)}), + ({}, {"cross_kv": torch.zeros(1)}), + ( + {}, + { + "mrope_rotary_cos_sin": torch.zeros(1, 8), + "attention_input_type": AttentionInputType.context_only, + }, + ), + ], +) +def test_gate_rejects_unsliced_side_inputs(metadata_extra, args_extra): + forward_args = AttentionForwardArgs(output=torch.empty(1), **args_extra) + with pytest.raises(NotImplementedError, match="does not support"): + trtllm_backend._locality_domain_attention_enabled( + _ATTN, _gate_metadata(**metadata_extra), forward_args + ) + + +def test_gate_rejects_sparse_indices(): + forward_args = AttentionForwardArgs(output=torch.empty(1)) + forward_args.sparse_runtime_params.sparse_attn_indices = torch.zeros(1, dtype=torch.int) + with pytest.raises(NotImplementedError, match="sparse"): + trtllm_backend._locality_domain_attention_enabled(_ATTN, _gate_metadata(), forward_args) + + +@pytest.mark.parametrize( + "kind, ld_slice", + [ + ("mixed", _slice(ctx_req=(0, 1), gen_req=(2, 4), ctx_tok=(0, 5), gen_tok=(9, 11), n=3)), + ("ctx_only", _slice(ctx_req=(1, 2), ctx_tok=(5, 9), n=1)), + ("gen_only", _slice(ctx_req=(2, 2), gen_req=(4, 6), gen_tok=(2, 4), n=2)), + ], +) +def test_scoped_metadata_counts_follow_the_slice(kind, ld_slice): + metadata = _Metadata(num_ctx=2, num_gen=4, ctx_tokens=(5, 4), gen_tokens=(2, 2)) + full = { + name: getattr(metadata, name) + for name in ("_num_contexts", "_num_generations", "_num_ctx_tokens", "_num_tokens") + } + with trtllm_backend._locality_domain_metadata(metadata, ld_slice, 1): + assert metadata.num_contexts == ld_slice["num_ctx_reqs"] + assert metadata.num_generations == ld_slice["num_gen_reqs"] + assert metadata._num_ctx_tokens == ld_slice["num_ctx_tokens"] + assert metadata._num_tokens == ld_slice["num_tokens"] + assert metadata.kv_lens_runtime is ld_slice["kv_lengths"] + assert metadata.block_ids_per_seq is None + for name, value in full.items(): + assert getattr(metadata, name) == value + assert metadata.block_ids_per_seq is not None + + +def test_scoped_metadata_restores_after_exception(): + metadata = _Metadata(num_ctx=1, num_gen=2, ctx_tokens=(3,), gen_tokens=(1, 1)) + ld_slice = _slice(ctx_req=(0, 1), gen_req=(1, 2), ctx_tok=(0, 3), gen_tok=(3, 4), n=2) + saved = dict(vars(metadata)) + with pytest.raises(ValueError): + with trtllm_backend._locality_domain_metadata(metadata, ld_slice, 0): + raise ValueError("injected") + for name, value in saved.items(): + assert getattr(metadata, name) is value or getattr(metadata, name) == value + + +def test_scoped_metadata_slices_generation_indexed_spec_inputs(): + num_ctx, num_gen = 1, 6 + metadata = _Metadata( + num_ctx, num_gen, ctx_tokens=(2,), gen_tokens=(3,) * 6, use_spec_decoding=True + ) + metadata.spec_decoding_generation_lengths = torch.arange(num_gen) + metadata.spec_decoding_packed_mask = torch.arange(num_gen * 6).view(num_gen, 3, 2) + metadata.spec_decoding_bl_tree_mask_offset = torch.arange(num_gen) * 10 + metadata.spec_bl_tree_first_sparse_mask_offset_kv = torch.arange(num_gen) * 100 + metadata.spec_decoding_bl_tree_mask = torch.zeros(16, dtype=torch.int) + metadata.spec_decoding_position_offsets = torch.arange(num_gen * 3) + metadata.spec_decoding_position_offsets_for_cpp = metadata.spec_decoding_position_offsets.view( + num_gen, 3 + ) + domain1_mask = torch.ones(16, dtype=torch.int) + metadata.locality_domain_spec_decoding_bl_tree_masks = [None, domain1_mask] + # Domain 1 owns generation requests 2..5 (batch rows 3..6). + ld_slice = _slice(ctx_req=(1, 1), gen_req=(3, 7), gen_tok=(6, 18), n=4) + with trtllm_backend._locality_domain_metadata(metadata, ld_slice, 1): + torch.testing.assert_close(metadata.spec_decoding_generation_lengths, torch.arange(2, 6)) + torch.testing.assert_close( + metadata.spec_decoding_packed_mask, torch.arange(12, 36).view(4, 3, 2) + ) + torch.testing.assert_close( + metadata.spec_decoding_bl_tree_mask_offset, torch.arange(2, 6) * 10 + ) + torch.testing.assert_close( + metadata.spec_bl_tree_first_sparse_mask_offset_kv, torch.arange(2, 6) * 100 + ) + torch.testing.assert_close( + metadata.spec_decoding_position_offsets, torch.arange(6, 18).view(4, 3) + ) + assert ( + metadata.spec_decoding_position_offsets_cpp is metadata.spec_decoding_position_offsets + ) + assert metadata.spec_decoding_bl_tree_mask is domain1_mask + assert metadata.spec_decoding_generation_lengths.numel() == num_gen + assert metadata.spec_decoding_position_offsets.dim() == 1 + assert metadata.spec_decoding_bl_tree_mask is not domain1_mask + + +def _coverage(prepared, kind, num_q_rows, metadata, num_kv_rows=None, dense_ctx_kv=False): + trtllm_backend._validate_locality_domain_coverage( + metadata, prepared, kind, num_q_rows, num_kv_rows, dense_ctx_kv + ) + + +def test_coverage_accepts_exact_tilings(): + metadata = _Metadata(num_ctx=2, num_gen=4, ctx_tokens=(5, 4), gen_tokens=(1,) * 4) + mixed = _prepared( + _slice(ctx_req=(0, 1), gen_req=(2, 4), ctx_tok=(0, 5), gen_tok=(9, 11)), + _slice(ctx_req=(1, 2), gen_req=(4, 6), ctx_tok=(5, 9), gen_tok=(11, 13)), + ) + _coverage(mixed, AttentionInputType.mixed, 13, metadata) + ctx_only = _prepared( + _slice(ctx_req=(0, 1), ctx_tok=(0, 5), ctx_kv=(0, 40)), + _slice(ctx_req=(1, 2), ctx_tok=(5, 9), ctx_kv=(40, 70)), + ) + _coverage( + ctx_only, AttentionInputType.context_only, 9, metadata, num_kv_rows=70, dense_ctx_kv=True + ) + gen_only = _prepared( + _slice(ctx_req=(1, 1), gen_req=(2, 4), gen_tok=(0, 2)), + _slice(ctx_req=(2, 2), gen_req=(4, 6), gen_tok=(2, 4)), + ) + _coverage(gen_only, AttentionInputType.generation_only, 4, metadata) + # One empty domain: the other must carry the whole batch. + all_in_one = _prepared(_slice(ctx_req=(0, 2), gen_req=(2, 6), ctx_tok=(0, 9), gen_tok=(9, 13))) + _coverage(all_in_one, AttentionInputType.mixed, 13, metadata) + + +@pytest.mark.parametrize( + "prepared, kind, num_q_rows, match", + [ + # Hole in the generation tokens: stale splits after a draft rewrite. + ( + _prepared( + _slice(ctx_req=(1, 1), gen_req=(2, 4), gen_tok=(0, 2)), + _slice(ctx_req=(2, 2), gen_req=(4, 6), gen_tok=(2, 4)), + ), + AttentionInputType.generation_only, + 12, + "generation tokens", + ), + # Overlapping context tokens. + ( + _prepared( + _slice(ctx_req=(0, 1), gen_req=(2, 4), ctx_tok=(0, 6), gen_tok=(9, 11)), + _slice(ctx_req=(1, 2), gen_req=(4, 6), ctx_tok=(5, 9), gen_tok=(11, 13)), + ), + AttentionInputType.mixed, + 13, + "context tokens", + ), + # Out of bounds request range. + ( + _prepared( + _slice(ctx_req=(0, 1), gen_req=(2, 4), ctx_tok=(0, 5), gen_tok=(9, 11)), + _slice(ctx_req=(1, 2), gen_req=(4, 7), ctx_tok=(5, 9), gen_tok=(11, 13)), + ), + AttentionInputType.mixed, + 13, + "generation requests", + ), + # Missing a domain entirely. + ( + _prepared(_slice(ctx_req=(0, 1), gen_req=(2, 4), ctx_tok=(0, 5), gen_tok=(9, 11))), + AttentionInputType.mixed, + 13, + "do not tile", + ), + # Negative start. + ( + _prepared(_slice(ctx_req=(0, 2), gen_req=(2, 6), ctx_tok=(0, 9), gen_tok=(-1, 13))), + AttentionInputType.mixed, + 13, + "generation tokens", + ), + ], +) +def test_coverage_rejects_partial_or_overlapping_slices(prepared, kind, num_q_rows, match): + metadata = _Metadata(num_ctx=2, num_gen=4, ctx_tokens=(5, 4), gen_tokens=(1,) * 4) + with pytest.raises(RuntimeError, match=match): + _coverage(prepared, kind, num_q_rows, metadata) + + +def test_coverage_rejects_token_major_presentation(): + metadata = _Metadata(num_ctx=0, num_gen=2, gen_tokens=(3, 3)) + metadata.kv_lens_runtime = torch.arange(6) + prepared = _prepared( + _slice(gen_req=(0, 1), gen_tok=(0, 3)), _slice(gen_req=(1, 2), gen_tok=(3, 6)) + ) + with pytest.raises(RuntimeError, match="rows for"): + _coverage(prepared, AttentionInputType.generation_only, 6, metadata) + + +def test_mrope_overrides_follow_kernel_indexing(): + num_ctx, num_q_rows, per_token = 2, 7, 4 + cos_sin = torch.arange(num_q_rows * per_token, dtype=torch.float).view(1, -1) + deltas = torch.arange(5).view(5, 1) + forward_args = AttentionForwardArgs( + output=torch.empty(1), mrope_rotary_cos_sin=cos_sin, mrope_position_deltas=deltas + ) + # Domain 1: context tokens [3, 7), generation requests 2..4 (rows 4..6). + ld_slice = _slice(ctx_req=(1, 2), gen_req=(4, 7), ctx_tok=(3, 7), gen_tok=(7, 10)) + overrides = trtllm_backend._locality_domain_mrope_overrides( + forward_args, ld_slice, num_ctx, num_q_rows + ) + torch.testing.assert_close( + overrides["mrope_rotary_cos_sin"], + torch.arange(3 * per_token, 7 * per_token, dtype=torch.float).view(1, -1), + ) + torch.testing.assert_close(overrides["mrope_position_deltas"], deltas[2:5]) + + # A domain with only generation rows gets no cos/sin; one with only + # context rows gets no deltas. + gen_only = trtllm_backend._locality_domain_mrope_overrides( + forward_args, _slice(gen_req=(2, 4), gen_tok=(0, 2)), num_ctx, num_q_rows + ) + assert gen_only["mrope_rotary_cos_sin"] is None + torch.testing.assert_close(gen_only["mrope_position_deltas"], deltas[0:2]) + ctx_only = trtllm_backend._locality_domain_mrope_overrides( + forward_args, _slice(ctx_req=(0, 1), ctx_tok=(0, 3)), num_ctx, num_q_rows + ) + assert ctx_only["mrope_position_deltas"] is None + + with pytest.raises(RuntimeError, match="cannot slice"): + trtllm_backend._locality_domain_mrope_overrides(forward_args, ld_slice, num_ctx, 5) + + +def test_draft_context_without_manager_disables_partitioning(): + metadata = SimpleNamespace(locality_domain_enabled=True) + with SpecWorkerBase.draft_kv_cache_context(None, metadata, None) as inner: + assert inner is metadata + assert metadata.locality_domain_enabled is False + assert metadata.locality_domain_enabled is True + + with pytest.raises(RuntimeError): + with SpecWorkerBase.draft_kv_cache_context(None, metadata, None): + raise RuntimeError("injected") + assert metadata.locality_domain_enabled is True + + with SpecWorkerBase.draft_kv_cache_context( + None, metadata, None, enable_draft_locality_domain_forkjoin=True + ): + assert metadata.locality_domain_enabled is True + + +def test_mla_gen_state_slices_per_domain_and_zeroes_separate_counters(monkeypatch): + ranges = {0: (0, 3), 1: (3, 5)} + monkeypatch.setattr( + trtllm_backend, + "_build_locality_domain_slices", + lambda metadata, ld, *args, **kwargs: _slice(gen_req=ranges[ld]), + ) + metadata = SimpleNamespace( + kv_cache_manager=SimpleNamespace(num_locality_domains=2), + num_contexts=0, + locality_domain_mla_gen_state=None, + ) + cu_q = torch.tensor([0, 1, 2, 3, 4, 5], dtype=torch.int32) + cu_kv = torch.tensor([0, 10, 20, 35, 40, 60], dtype=torch.int32) + counter = torch.ones(1, dtype=torch.int32) + + trtllm_backend._prepare_locality_domain_mla_gen_state_from_full( + metadata, cu_q, cu_kv, counter, None, None, total_gen_tokens=5 + ) + + first, second = metadata.locality_domain_mla_gen_state + assert first["cu_q_seqlens"].tolist() == [0, 1, 2, 3] + assert first["cu_kv_seqlens"].tolist() == [0, 10, 20, 35] + assert second["cu_q_seqlens"].tolist() == [0, 1, 2] + assert second["cu_kv_seqlens"].tolist() == [0, 5, 25] + # A slice from row 0 is a view; only the later domain needs a rebased copy. + assert first["cu_q_seqlens"].data_ptr() == cu_q.data_ptr() + assert second["cu_q_seqlens"].data_ptr() != cu_q.data_ptr() + assert first["fmha_scheduler_counter"].tolist() == [0] + assert second["fmha_scheduler_counter"].tolist() == [0] + assert first["fmha_scheduler_counter"].data_ptr() != second["fmha_scheduler_counter"].data_ptr() + assert counter.tolist() == [1] diff --git a/tests/unittest/_torch/executor/test_context_only_warmup.py b/tests/unittest/_torch/executor/test_context_only_warmup.py index e7c5c64cee37..d0574f325421 100644 --- a/tests/unittest/_torch/executor/test_context_only_warmup.py +++ b/tests/unittest/_torch/executor/test_context_only_warmup.py @@ -23,6 +23,7 @@ def _warmup_fixture(context_only: bool, available: int = 64): context_swa_layer_limit=20 if context_only else None, tokens_per_block=128, num_extra_kv_tokens=0, + num_locality_domains=1, _has_cp_helix=False, _stream=SimpleNamespace(cuda_stream=0), _create_kv_cache=Mock( diff --git a/tests/unittest/_torch/executor/test_locality_domain_dispatch.py b/tests/unittest/_torch/executor/test_locality_domain_dispatch.py new file mode 100644 index 000000000000..1ee8f9c08dbc --- /dev/null +++ b/tests/unittest/_torch/executor/test_locality_domain_dispatch.py @@ -0,0 +1,114 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Locality domain split arithmetic of the decoder runner, without a model.""" + +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.locality_domain import ( + balanced_generation_dispatch, + locality_domain_row_splits, + sort_generation_requests_by_locality_domain, +) + +pytestmark = pytest.mark.cpu_only + + +def _req(request_id, domain, dummy=False, batch_idx=0): + return SimpleNamespace( + py_request_id=request_id, + py_locality_domain_id=domain, + is_cuda_graph_dummy=dummy, + is_dummy=dummy, + py_batch_idx=None if dummy else batch_idx, + ) + + +@pytest.mark.parametrize("num_rows", range(0, 9)) +def test_balanced_dispatch_depends_only_on_row_count(num_rows): + dispatch = balanced_generation_dispatch(num_rows, 2) + assert len(dispatch) == num_rows + assert dispatch == sorted(dispatch) + assert dispatch.count(0) == (num_rows + 1) // 2 + assert dispatch.count(1) == num_rows // 2 + + +def test_balanced_dispatch_spreads_remainder_over_leading_domains(): + assert balanced_generation_dispatch(7, 3) == [0, 0, 0, 1, 1, 2, 2] + + +@pytest.mark.parametrize("num_real", [1, 2, 3]) +@pytest.mark.parametrize("padding", [0, 1, 5, 6]) +def test_repeated_padding_dummies_count_per_occurrence(num_real, padding): + dummy = _req(-1, 0, dummy=True) + packed = [_req(i, i % 2) for i in range(num_real)] + [dummy] * padding + reqs, tokens = locality_domain_row_splits(packed, [1] * len(packed), 2, positional=True) + total = num_real + padding + assert reqs == [(total + 1) // 2, total // 2] + assert tokens == reqs + + +def test_same_padded_size_gives_same_split_for_any_composition(): + dummy = _req(-1, 1, dummy=True) + a = [_req(0, 0), _req(1, 1)] + [dummy] * 6 + b = [_req(i, 1) for i in range(5)] + [dummy] * 3 + assert locality_domain_row_splits(a, [1] * 8, 2, positional=True) == ( + locality_domain_row_splits(b, [1] * 8, 2, positional=True) + ) + + +def test_ownership_split_counts_tokens_per_domain(): + packed = [_req(0, 0), _req(1, 0), _req(2, 1)] + reqs, tokens = locality_domain_row_splits(packed, [3, 1, 2], 2, positional=False) + assert reqs == [2, 1] + assert tokens == [4, 2] + + +def test_ownership_split_falls_back_to_positional_when_interleaved(): + packed = [_req(0, 0), _req(1, 1), _req(2, 0)] + assert locality_domain_row_splits(packed, [1, 1, 1], 2, positional=False) == ( + locality_domain_row_splits(packed, [1, 1, 1], 2, positional=True) + ) + + +def test_ownership_split_falls_back_to_positional_when_domains_descend(): + # Contiguous but out of domain order (an extend row before a generation + # row): ownership counts [1, 4] would hand domain 0 the first token of the + # domain-1 request, so the split must be positional. + packed = [_req(0, 1), _req(1, 0)] + assert locality_domain_row_splits(packed, [4, 1], 2, positional=False) == ([1, 1], [4, 1]) + + +def test_sort_keeps_padding_dummies_last(): + dummy = _req(-1, 0, dummy=True) + reqs = [dummy, _req(0, 1), _req(1, 0), dummy] + ordered = sort_generation_requests_by_locality_domain(reqs) + assert [r.py_request_id for r in ordered] == [1, 0, -1, -1] + + +def test_sort_keeps_host_token_rows_ahead_of_previous_step_rows(): + # The overlap scheduler packs host tokens before the previous step's device + # tokens; a domain-only sort would hand row 2's host token to row 0. + reqs = [ + _req(0, 0), + _req(1, 1, batch_idx=None), + _req(2, 0, batch_idx=None), + _req(3, 1), + _req(4, 0), + _req(-1, 0, dummy=True), + ] + ordered = sort_generation_requests_by_locality_domain(reqs, frozenset({4})) + assert [r.py_request_id for r in ordered] == [2, 4, 1, 0, 3, -1] diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 06bd4e3ddb1f..8a81be6a5060 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -6,7 +6,7 @@ from contextlib import nullcontext from dataclasses import dataclass, replace from types import SimpleNamespace -from unittest.mock import Mock, patch +from unittest.mock import Mock, PropertyMock, patch import pytest import torch @@ -2753,6 +2753,59 @@ def test_pad_batch_strips_cudagraph_dummies_on_clean_exit(self) -> None: kv_cache_manager.shutdown() + def _warmup_request_split(self, num_locality_domains): + model_engine, kv_cache_manager = create_model_engine_and_kvcache() + resource_manager = Mock() + resource_manager.get_resource_manager.side_effect = ( + lambda key: kv_cache_manager + if key == model_engine.kv_cache_manager_key else None) + + kv_cache_manager.get_num_free_blocks = Mock(return_value=8) + kv_cache_manager.max_seq_len = 2048 + kv_cache_manager.get_num_available_tokens = Mock(return_value=100) + + # update_draft_len() runs over the returned requests: it measures + # py_draft_tokens and or-assigns py_needs_onehot_draft_probs, so the + # stand-ins need real values rather than auto-created Mock attributes. + def _dummy_request(): + return Mock(py_draft_tokens=[], py_needs_onehot_draft_probs=False) + + num_long = min(num_locality_domains, 4) + kv_cache_manager.add_dummy_requests = Mock(side_effect=[ + [_dummy_request() for _ in range(4 - num_long)], + [_dummy_request() for _ in range(num_long)], + ]) + try: + with patch.object(KVCacheManagerV2, + "num_locality_domains", + new_callable=PropertyMock, + return_value=num_locality_domains): + model_engine._runner._dummy_requests.create_cuda_graph_warmup_request( + resource_manager, + batch_size=4, + draft_len=0, + max_seq_len=512) + return kv_cache_manager.add_dummy_requests.call_args_list + finally: + kv_cache_manager.shutdown() + + def test_cuda_graph_warmup_request_split(self): + short_call, long_call = self._warmup_request_split(1) + self.assertEqual(short_call.args[0], [0, 1, 2]) + self.assertNotIn("locality_domain_ids", short_call.kwargs) + self.assertEqual(long_call.kwargs["request_ids"], [3]) + self.assertEqual(long_call.kwargs["token_nums"], [100]) + self.assertNotIn("locality_domain_ids", long_call.kwargs) + + def test_cuda_graph_warmup_request_split_per_locality_domain(self): + # One longest request per domain, short requests dealt from domain 0. + short_call, long_call = self._warmup_request_split(2) + self.assertEqual(short_call.args[0], [0, 1]) + self.assertEqual(short_call.kwargs["locality_domain_ids"], [0, 1]) + self.assertEqual(long_call.kwargs["request_ids"], [2, 3]) + self.assertEqual(long_call.kwargs["token_nums"], [100, 100]) + self.assertEqual(long_call.kwargs["locality_domain_ids"], [0, 1]) + def test_pad_batch_strips_cudagraph_dummies_on_exception(self) -> None: # The strip must fire even when the body raises. This is the # critical property of `finally` vs. a plain trailing statement — From 5bccaf683e9bc05431725bb237f24827d57adf5d Mon Sep 17 00:00:00 2001 From: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> Date: Fri, 23 Oct 2026 18:38:50 +0000 Subject: [PATCH 4/5] [None][perf] Size the MLA decode kernel's KV split for the SMs of its locality domain A kernel launched on a locality-domain stream runs only on that Green Context's SMs, but MLA decode picked the trtllm-gen kernel and sized its KV split for the whole device, launching about twice the CTAs one wave could hold. The localization handle now records the SM count of every stream it creates, and MLA generation passes it to the FMHA runner in place of the device count. The lookup is a table read, so it is safe during CUDA graph capture. GQA decode keeps the device count: at batch 1-16 its kernels ran 5-10% slower with it. Signed-off-by: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> --- cpp/tensorrt_llm/common/attentionOp.cpp | 7 ++- .../locality_domain/locality_domain_utils.cpp | 59 +++++++++++++++++++ .../locality_domain/locality_domain_utils.h | 4 ++ .../unit_tests/runtime/localizationTest.cu | 45 ++++++++++++++ 4 files changed, 114 insertions(+), 1 deletion(-) diff --git a/cpp/tensorrt_llm/common/attentionOp.cpp b/cpp/tensorrt_llm/common/attentionOp.cpp index 7d0e74b5ef06..0a2ba65f25d2 100644 --- a/cpp/tensorrt_llm/common/attentionOp.cpp +++ b/cpp/tensorrt_llm/common/attentionOp.cpp @@ -32,6 +32,7 @@ #include "tensorrt_llm/kernels/sparseAttentionKernels.h" #include "tensorrt_llm/kernels/unfusedAttentionKernels.h" #include "tensorrt_llm/runtime/iBuffer.h" +#include "tensorrt_llm/runtime/locality_domain/locality_domain_utils.h" #include "tensorrt_llm/runtime/utils/debugUtils.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" #include @@ -1221,7 +1222,11 @@ int AttentionOp::mlaGeneration( // Set it to INT_MAX as the kv cache pageOffsets will ensure that there is no out-of-bounds access. tllmRunnerParams.mNumPagesInMemPool = INT_MAX; - tllmRunnerParams.mMultiProcessorCount = mMultiProcessorCount; + // A locality-domain stream runs only on its Green Context's SMs, so size the KV split for those. + // GQA decode stays on the device count: at small batch its kernels run faster over-split. + auto const streamSmCount = static_cast(tensorrt_llm::locality_domain::getLocalizedStreamSmCount(stream)); + tllmRunnerParams.mMultiProcessorCount + = streamSmCount > 0 ? std::min(streamSmCount, mMultiProcessorCount) : mMultiProcessorCount; tllmRunnerParams.stream = stream; tllmRunnerParams.mSfStartTokenIdx = generation_params.start_token_idx_sf; tllmRunnerParams.mSkipCorrThreshold = mSkipCorrectionThreshold; diff --git a/cpp/tensorrt_llm/runtime/locality_domain/locality_domain_utils.cpp b/cpp/tensorrt_llm/runtime/locality_domain/locality_domain_utils.cpp index 011fe9bd5f02..0ebf8d5e61b0 100644 --- a/cpp/tensorrt_llm/runtime/locality_domain/locality_domain_utils.cpp +++ b/cpp/tensorrt_llm/runtime/locality_domain/locality_domain_utils.cpp @@ -31,6 +31,7 @@ #include #include #include +#include #include namespace tensorrt_llm::locality_domain @@ -39,6 +40,41 @@ namespace tensorrt_llm::locality_domain namespace { +// SM count of every stream this file creates on a Green Context. Kernel launchers read it while a +// CUDA graph is being captured, where querying the stream's resources is not allowed. +class LocalizedStreamSmCounts +{ +public: + static LocalizedStreamSmCounts& instance() + { + static auto* const counts = new LocalizedStreamSmCounts; + return *counts; + } + + void add(CUstream stream, unsigned int smCount) + { + std::lock_guard const lock{mMutex}; + mCounts[stream] = smCount; + } + + void remove(CUstream stream) + { + std::lock_guard const lock{mMutex}; + mCounts.erase(stream); + } + + unsigned int find(CUstream stream) const + { + std::lock_guard const lock{mMutex}; + auto const it = mCounts.find(stream); + return it == mCounts.end() ? 0 : it->second; + } + +private: + mutable std::mutex mMutex; + std::map mCounts; +}; + #if CUDA_VERSION >= 13040 constexpr int kLocalityDomainApiVersion = 13'040; constexpr int kLocalityDomainCount = detail::kLocalityDomainCount; @@ -385,6 +421,10 @@ class GreenContextPartitions result = mApi.greenCtxStreamCreate( &mRemainderStream, mRemainderGreenContext, CU_STREAM_NON_BLOCKING, /*priority=*/0); } + if (result == CUDA_SUCCESS) + { + LocalizedStreamSmCounts::instance().add(mRemainderStream, getRemainderSmCount()); + } if (result != CUDA_SUCCESS) { TLLM_LOG_WARNING( @@ -425,6 +465,7 @@ class GreenContextPartitions return result; } mLocalizedStreams[localityDomainId] = newStream; + LocalizedStreamSmCounts::instance().add(newStream, mLocalizedResources[localityDomainId].sm.smCount); } *stream = mLocalizedStreams[localityDomainId]; return CUDA_SUCCESS; @@ -476,6 +517,7 @@ class GreenContextPartitions { if (stream != nullptr) { + LocalizedStreamSmCounts::instance().remove(stream); CUresult const result = cuStreamDestroy(stream); if (result != CUDA_SUCCESS) { @@ -488,6 +530,7 @@ class GreenContextPartitions if (mRemainderStream != nullptr) { + LocalizedStreamSmCounts::instance().remove(mRemainderStream); CUresult const result = cuStreamDestroy(mRemainderStream); if (result != CUDA_SUCCESS) { @@ -1114,6 +1157,22 @@ CUstream LocalizationHandle::getReservedRemainderStream() return mImpl->getReservedRemainderStream(); } +unsigned int getLocalizedStreamSmCount(CUstream stream) noexcept +{ + if (stream == nullptr) + { + return 0; + } + try + { + return LocalizedStreamSmCounts::instance().find(stream); + } + catch (std::system_error const&) + { + return 0; + } +} + bool deviceSupportsLocalization(int device) noexcept { #if CUDA_VERSION >= 13040 diff --git a/cpp/tensorrt_llm/runtime/locality_domain/locality_domain_utils.h b/cpp/tensorrt_llm/runtime/locality_domain/locality_domain_utils.h index e57aa238d9f5..b51705d0c9f8 100644 --- a/cpp/tensorrt_llm/runtime/locality_domain/locality_domain_utils.h +++ b/cpp/tensorrt_llm/runtime/locality_domain/locality_domain_utils.h @@ -78,6 +78,10 @@ class LocalizationHandle Localization* mImpl; }; +//! SM count of a stream created by LocalizationHandle (a locality-domain or remainder stream), or 0 for any +//! other stream. A table lookup with no CUDA call, so it is safe during CUDA graph capture. +unsigned int getLocalizedStreamSmCount(CUstream stream) noexcept; + bool deviceSupportsLocalization(int device) noexcept; } // namespace locality_domain diff --git a/cpp/tests/unit_tests/runtime/localizationTest.cu b/cpp/tests/unit_tests/runtime/localizationTest.cu index 75c56d38d0a5..287d76acc26d 100644 --- a/cpp/tests/unit_tests/runtime/localizationTest.cu +++ b/cpp/tests/unit_tests/runtime/localizationTest.cu @@ -441,6 +441,51 @@ TEST_F(LocalizationTest, CreateLocalizedStreamLocalityDomain1) EXPECT_EQ(mHandle->createLocalizedStream(1), stream); } +TEST_F(LocalizationTest, LocalizedStreamSmCountMatchesItsPartition) +{ + if (!mSupportsLocalityDomain) + { + GTEST_SKIP() << "LOCALITY_DOMAIN not supported, skipping test"; + } + + for (int localityDomainId = 0; localityDomainId < 2; ++localityDomainId) + { + CUstream stream = mHandle->createLocalizedStream(localityDomainId); + auto const [partitionSmCount, deviceSmCount] = mHandle->getLocalityDomainComputeSmCounts(localityDomainId); + ASSERT_GT(partitionSmCount, 0u); + EXPECT_LT(partitionSmCount, deviceSmCount); + EXPECT_EQ(getLocalizedStreamSmCount(stream), partitionSmCount); + } +} + +TEST_F(LocalizationTest, OrdinaryStreamHasNoLocalizedSmCount) +{ + EXPECT_EQ(getLocalizedStreamSmCount(nullptr), 0u); + + cudaStream_t stream{}; + ASSERT_EQ(cudaStreamCreate(&stream), cudaSuccess); + EXPECT_EQ(getLocalizedStreamSmCount(reinterpret_cast(stream)), 0u); + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + +TEST_F(LocalizationTest, LocalizedStreamSmCountIsReadableDuringGraphCapture) +{ + if (!mSupportsLocalityDomain) + { + GTEST_SKIP() << "LOCALITY_DOMAIN not supported, skipping test"; + } + + CUstream stream = mHandle->createLocalizedStream(0); + auto const expected = mHandle->getLocalityDomainComputeSmCounts(0).first; + auto const cudaStream = reinterpret_cast(stream); + ASSERT_EQ(cudaStreamBeginCapture(cudaStream, cudaStreamCaptureModeThreadLocal), cudaSuccess); + unsigned int const captured = getLocalizedStreamSmCount(stream); + cudaGraph_t graph{}; + ASSERT_EQ(cudaStreamEndCapture(cudaStream, &graph), cudaSuccess); + ASSERT_EQ(cudaGraphDestroy(graph), cudaSuccess); + EXPECT_EQ(captured, expected); +} + // Bandwidth checks allocate tens of GiB and are disabled in normal unit-test runs. // Run them explicitly with --gtest_also_run_disabled_tests when benchmarking locality-domain locality. From 5c2badddf01d99cc77564de347828572ee20f862 Mon Sep 17 00:00:00 2001 From: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> Date: Fri, 23 Oct 2026 18:38:52 +0000 Subject: [PATCH 5/5] [None][feat] Transfer localized KV pools in disaggregated serving A localized pool is backed only on its domains' mapped ranges, so the transceiver registers each backed range of a pool rather than the whole extent, and the page table carries the ranges to the peer. The single-process transceiver test transfers between localized and non-localized pools, MHA and MLA, at TP 1->1, 2->1 and 1->2, with a head size that leaves a hole in the pool. Signed-off-by: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> --- .../disaggregation/resource/kv_extractor.py | 13 +++ .../_torch/disaggregation/resource/page.py | 13 ++- .../_torch/disaggregation/resource/utils.py | 31 ++++--- .../disaggregated/region/test_page.py | 64 ++++++++++++++ .../test_cache_transceiver_single_process.py | 85 +++++++++++++++++-- .../unittest/disaggregated/test_extractor.py | 43 +++++++++- 6 files changed, 229 insertions(+), 20 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py index b1132c10ddc4..e3c4bd414710 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py +++ b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py @@ -327,6 +327,14 @@ def _build_pool_views_for_variant( return pool_views +def _mapped_ranges_of(pool) -> Optional[Tuple[Tuple[int, int], ...]]: + """Mapped byte ranges of a V2 pool, or None when its whole extent is backed.""" + ranges = getattr(pool, "mapped_ranges", None) + if not ranges: + return None + return tuple((int(start), int(end)) for start, end in ranges) + + def _build_page_table_v2(manager) -> KVCachePageTable: """Build a KVCachePageTable from a KVCacheManagerV2. @@ -408,6 +416,11 @@ def _window_size_for_layer(internal_layer_id: int): base_address=int(pool.base_address), slot_bytes=int(pool.slot_bytes), num_slots=int(pg_desc.num_slots), + mapped_ranges=( + _mapped_ranges_of(pool) + if getattr(manager, "num_locality_domains", 1) > 1 + else None + ), ) for pool in pg_desc.pools ] diff --git a/tensorrt_llm/_torch/disaggregation/resource/page.py b/tensorrt_llm/_torch/disaggregation/resource/page.py index 1254c06001f6..d376f63d6210 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/page.py +++ b/tensorrt_llm/_torch/disaggregation/resource/page.py @@ -109,6 +109,9 @@ class PhysicalPool: slot_bytes: int num_slots: int slot_stride_bytes: Optional[int] = None + # Byte ranges relative to base_address that are backed by memory; None when the whole + # extent is. A localized KV pool maps each locality domain's chunks at distinct offsets. + mapped_ranges: Optional[Tuple[Tuple[int, int], ...]] = None def __post_init__(self) -> None: if self.slot_stride_bytes is None: @@ -117,12 +120,15 @@ def __post_init__(self) -> None: raise ValueError("slot_stride_bytes must be greater than or equal to slot_bytes") def to_dict(self) -> dict: - return { + result = { "base_address": int(self.base_address), "slot_bytes": int(self.slot_bytes), "num_slots": int(self.num_slots), "slot_stride_bytes": int(self.slot_stride_bytes), } + if self.mapped_ranges is not None: + result["mapped_ranges"] = [list(bounds) for bounds in self.mapped_ranges] + return result @staticmethod def from_dict(data: dict) -> "PhysicalPool": @@ -135,6 +141,11 @@ def from_dict(data: dict) -> "PhysicalPool": if data.get("slot_stride_bytes") is not None else None ), + mapped_ranges=( + tuple((int(start), int(end)) for start, end in data["mapped_ranges"]) + if data.get("mapped_ranges") is not None + else None + ), ) diff --git a/tensorrt_llm/_torch/disaggregation/resource/utils.py b/tensorrt_llm/_torch/disaggregation/resource/utils.py index 579cc42d1d13..43d0b144c91d 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/utils.py +++ b/tensorrt_llm/_torch/disaggregation/resource/utils.py @@ -243,7 +243,19 @@ def get_unique_pool_memory_descs( ) -> list[tuple[int, int, int, str]]: """Return deduplicated (ptr, size, device_id, name) tuples for all physical pools.""" unique_pools: dict[tuple[int, int], int] = {} # (ptr, size) -> index - pool_counter = 0 + + def add_pool(pool: PhysicalPool, size: int) -> None: + # A pool with holes registers each backed range on its own; a region already inside a + # registered one is skipped. + ranges = pool.mapped_ranges if pool.mapped_ranges is not None else ((0, size),) + for start, end in ranges: + base = pool.base_address + start + if end <= start or any( + b <= base and base + end - start <= b + s for b, s in unique_pools + ): + continue + unique_pools[(base, end - start)] = len(unique_pools) + for lg_idx, lg in enumerate(page_table.layer_groups): if lg.kind == CacheKind.STATE: assert isinstance(lg, MambaLayerGroup) @@ -256,21 +268,18 @@ def get_unique_pool_memory_descs( key=lambda p: p.base_address, ) for pool in pools: - pool_size = pool.num_slots * pool.slot_stride_bytes - if any(b <= pool.base_address < b + s for (b, s) in unique_pools): + # Interleaved role views start inside the same allocation; + # their nominal stride extent can run past its end. + if pool.mapped_ranges is None and any( + b <= pool.base_address < b + s for b, s in unique_pools + ): continue - pool_key = (pool.base_address, pool_size) - if pool_key not in unique_pools: - unique_pools[pool_key] = pool_counter - pool_counter += 1 + add_pool(pool, pool.num_slots * pool.slot_stride_bytes) else: # PAGED (attention): each pool view is an independent allocation for pv in lg.pool_views: pool = get_physical_pool(page_table, lg_idx, pv.pool_idx) - pool_key = (pool.base_address, pool.num_slots * pool.slot_stride_bytes) - if pool_key not in unique_pools: - unique_pools[pool_key] = pool_counter - pool_counter += 1 + add_pool(pool, pool.num_slots * pool.slot_stride_bytes) return [ (pool_ptr, pool_size, device_id, f"kv_cache_memory_pool{idx}") for (pool_ptr, pool_size), idx in unique_pools.items() diff --git a/tests/unittest/disaggregated/region/test_page.py b/tests/unittest/disaggregated/region/test_page.py index b8704ae78382..43d839386c9f 100644 --- a/tests/unittest/disaggregated/region/test_page.py +++ b/tests/unittest/disaggregated/region/test_page.py @@ -263,3 +263,67 @@ def test_layer_byte_ranges_empty_view_raises(): with pytest.raises(ValueError, match="no buffer entries"): get_layer_byte_ranges(_view([])) + + +@pytest.mark.parametrize("ranges", [None, (), ((0, 512), (1024, 2048))]) +def test_physical_pool_mapped_ranges_roundtrip(ranges): + import json + + pool = PhysicalPool(0x10000, 256, 8, mapped_ranges=ranges) + wire = json.loads(json.dumps(pool.to_dict())) + assert PhysicalPool.from_dict(wire).mapped_ranges == ranges + if ranges is None: + assert "mapped_ranges" not in wire + + +def test_registration_excludes_holes_and_deduplicates_shared_pools(): + from tensorrt_llm._torch.disaggregation.resource.utils import get_unique_pool_memory_descs + + pool = PhysicalPool(0x10000, 256, 8, mapped_ranges=((0, 512), (1024, 2048))) + table = KVCachePageTable( + tokens_per_block=64, + layer_groups=[ + AttentionLayerGroup( + pool_group_idx=0, + kv_head_num_per_rank=8, + sliding_window_size=None, + local_layers=[LocalLayer(0, 5)], + pool_views=[PoolView(pool_idx=0, buffer_entries=_make_buffer_entries())], + ) + ] + * 2, + pool_groups=[PhysicalPoolGroup(pools=[pool])], + ) + assert get_unique_pool_memory_descs(table, device_id=3) == [ + (0x10000, 512, 3, "kv_cache_memory_pool0"), + (0x10400, 1024, 3, "kv_cache_memory_pool1"), + ] + + +def test_registration_preserves_interleaved_state_pool_deduplication(): + from tensorrt_llm._torch.disaggregation.resource.page import MambaLayerGroup + from tensorrt_llm._torch.disaggregation.resource.utils import get_unique_pool_memory_descs + + table = KVCachePageTable( + tokens_per_block=64, + layer_groups=[ + MambaLayerGroup( + pool_group_idx=0, + local_layers=[LocalLayer(0, 0)], + pool_views=[ + PoolView(pool_idx=i, buffer_entries=_make_buffer_entries()) for i in (1, 0) + ], + ) + ], + pool_groups=[ + PhysicalPoolGroup( + pools=[ + PhysicalPool(0x10000, 128, 4, slot_stride_bytes=256), + PhysicalPool(0x10080, 128, 4, slot_stride_bytes=256), + ] + ) + ], + ) + assert get_unique_pool_memory_descs(table, device_id=0) == [ + (0x10000, 1024, 0, "kv_cache_memory_pool0") + ] diff --git a/tests/unittest/disaggregated/test_cache_transceiver_single_process.py b/tests/unittest/disaggregated/test_cache_transceiver_single_process.py index e6b39fafccd3..4ab7081e4545 100644 --- a/tests/unittest/disaggregated/test_cache_transceiver_single_process.py +++ b/tests/unittest/disaggregated/test_cache_transceiver_single_process.py @@ -334,6 +334,7 @@ def _create_cache_manager( max_attention_window_vec: Optional[List[int]] = None, num_layers: int = NUM_LAYERS, max_batch_size: int = MAX_BATCH_SIZE, + enable_locality_domains: bool = False, ) -> "KVCacheManagerV2": """Create a KVCacheManagerV2 for the given mapping.""" num_kv_heads = 1 if is_mla else NUM_KV_HEADS @@ -358,6 +359,7 @@ def _create_cache_manager( mapping=mapping, dtype=DataType.FLOAT, vocab_size=VOCAB_SIZE, + enable_locality_domains=enable_locality_domains, ) @@ -371,6 +373,7 @@ def _create_managers_for_instance( num_layers: int = NUM_LAYERS, max_batch_size: int = MAX_BATCH_SIZE, cp_size: int = 1, + enable_locality_domains: bool = False, ) -> list[KVCacheManagerV2]: """Create cache managers for all ranks in an instance. @@ -402,6 +405,7 @@ def _create_managers_for_instance( max_attention_window_vec, num_layers, max_batch_size, + enable_locality_domains, ) ) return managers @@ -426,16 +430,19 @@ def _init_pool_data_v2( tp_rank = rank % tp page_table = KVRegionExtractorV1(mgr).page_table - unique_pools: Dict[int, int] = {} + # A localized pool leaves unmapped holes between its domains' chunks, so only + # the backed ranges are written. + unique_ranges: Dict[int, int] = {} for pg in page_table.pool_groups: for pool_desc in pg.pools: - key = pool_desc.base_address pool_bytes = pool_desc.slot_bytes * pool_desc.num_slots - if key not in unique_pools or pool_bytes > unique_pools[key]: - unique_pools[key] = pool_bytes + for start, end in pool_desc.mapped_ranges or ((0, pool_bytes),): + key = pool_desc.base_address + start + if key not in unique_ranges or end - start > unique_ranges[key]: + unique_ranges[key] = end - start element_bytes = get_size_in_bytes(1, mgr.dtype) - for pool_base_ptr, pool_size in unique_pools.items(): + for pool_base_ptr, pool_size in unique_ranges.items(): pool_size_elements = pool_size // element_bytes pool_tensor = convert_to_torch_tensor( TensorWrapper(pool_base_ptr, mgr.dtype, [pool_size_elements]) @@ -524,7 +531,9 @@ def _add_sequence( harness cannot (it transfers right after a prompt-only resize, so valid_only skips the page-less tail). Used to exercise the SWA + MTP transfer path. """ - kv_cache = mgr._create_kv_cache(request_id, None, None) + kv_cache = mgr._create_kv_cache( + request_id, None, None, locality_domain_id=mgr.pick_locality_domain(request_id) + ) if not mgr.enable_block_reuse: kv_cache.stop_committing() success = kv_cache.resume(torch.cuda.current_stream().cuda_stream) @@ -967,6 +976,8 @@ def run_transfer_test( expect_cpp_bounce: bool = False, gen_scratch_tokens: int = 0, gen_cp: int = 1, + ctx_locality_domains: bool = False, + gen_locality_domains: bool = False, ) -> None: """Run a full KV transfer test using KvCacheTransceiverV2. @@ -984,6 +995,10 @@ def run_transfer_test( incompatible with gen_scratch_tokens > 0. HELIX VALIDATE: the whole gen_cp path (dist-mock collectives, ledger block ownership, strided verify) is exercised only on GPU -- validate end to end. + + ctx_locality_domains / gen_locality_domains build that side's managers with + locality domains; under TRT_LLM_MOCK_LOCALIZATION_SUPPORT=1 their pools are split + into per-domain chunks with unmapped holes on any GPU. """ if request_lengths is None: request_lengths = REQUEST_LENGTHS @@ -1004,6 +1019,7 @@ def run_transfer_test( max_attention_window_vec, num_layers, max_batch_size, + enable_locality_domains=ctx_locality_domains, ) gen_managers = _create_managers_for_instance( gen_tp, @@ -1015,6 +1031,7 @@ def run_transfer_test( num_layers, max_batch_size, cp_size=gen_cp, + enable_locality_domains=gen_locality_domains, ) # 2. Initialize data: random for ctx, zeros for gen @@ -1474,6 +1491,62 @@ def test_cache_transceiver_uneven_pp( print("PASSED") +def _pool_has_holes(manager: KVCacheManagerV2) -> bool: + for pg in KVRegionExtractorV1(manager).page_table.pool_groups: + for pool in pg.pools: + if pool.mapped_ranges is None: + continue + covered = sum(end - start for start, end in pool.mapped_ranges) + if covered < pool.slot_bytes * pool.num_slots: + return True + return False + + +@pytest.mark.timeout(180) +@pytest.mark.parametrize( + "ctx_tp,gen_tp", [(1, 1), (2, 1), (1, 2)], ids=["tp1_tp1", "tp2_tp1", "tp1_tp2"] +) +@pytest.mark.parametrize( + "ctx_locality_domains,gen_locality_domains", + [(True, True), (True, False), (False, True)], + ids=["ld_both", "ld_ctx", "ld_gen"], +) +@pytest.mark.parametrize("is_mla", [False, True], ids=["mha", "mla"]) +def test_cache_transceiver_locality_domains( + monkeypatch, ctx_tp, gen_tp, ctx_locality_domains, gen_locality_domains, is_mla +): + """Transfer between localized and non-localized KV pools. + + A localized pool maps each domain's chunks at separate offsets of one VA + reservation, so registration and the pool fill must cover the backed ranges + only; the round-trip data check proves both sides addressed the same pages. + """ + monkeypatch.setenv("TRT_LLM_MOCK_LOCALIZATION_SUPPORT", "1") + # A 72-wide head gives slot sizes whose domain chunks span several granules, which + # leaves an unmapped hole between the domains (power-of-two heads pack them flush). + monkeypatch.setitem(globals(), "HEAD_DIM", 72) + probe = _create_cache_manager( + Mapping(world_size=1, rank=0, tp_size=1), is_mla, True, enable_locality_domains=True + ) + try: + assert probe.num_locality_domains == 2 + assert _pool_has_holes(probe) + finally: + probe.shutdown() + run_transfer_test( + ctx_tp=ctx_tp, + ctx_pp=1, + gen_tp=gen_tp, + gen_pp=1, + ctx_enable_dp=False, + gen_enable_dp=False, + is_mla=is_mla, + use_v2=True, + ctx_locality_domains=ctx_locality_domains, + gen_locality_domains=gen_locality_domains, + ) + + @pytest.mark.timeout(180) @pytest.mark.parametrize("is_mla", [False, True], ids=["mha", "mla"]) @pytest.mark.parametrize("gen_cp", [2, 4], ids=["cp2", "cp4"]) diff --git a/tests/unittest/disaggregated/test_extractor.py b/tests/unittest/disaggregated/test_extractor.py index 7e7e3f7f0803..b9c22860f8bb 100644 --- a/tests/unittest/disaggregated/test_extractor.py +++ b/tests/unittest/disaggregated/test_extractor.py @@ -252,7 +252,15 @@ def test_layer_group_meta_serialization(): assert len(restored_lg.pool_views[0].buffer_entries) == 2 -def _make_fake_v2_manager(attrs, role_mapper_kinds, *, num_pools=1, slot_bytes_list=(640,)): +def _make_fake_v2_manager( + attrs, + role_mapper_kinds, + *, + num_pools=1, + slot_bytes_list=(640,), + num_locality_domains=1, + mapped_ranges=None, +): """Build a minimal duck-typed KVCacheManagerV2 for _build_page_table_v2. ``attrs`` keeps the storage-layer shape ``{(layer_id, role): @@ -290,7 +298,9 @@ def _make_fake_v2_manager(attrs, role_mapper_kinds, *, num_pools=1, slot_bytes_l pool_group_index=0, num_slots=4, pools=[ - SimpleNamespace(pool_index=pi, base_address=base, slot_bytes=sb) + SimpleNamespace( + pool_index=pi, base_address=base, slot_bytes=sb, mapped_ranges=mapped_ranges + ) for pi, sb in enumerate(slot_bytes_list) ], slot_desc=SimpleNamespace( @@ -313,11 +323,40 @@ def _make_fake_v2_manager(attrs, role_mapper_kinds, *, num_pools=1, slot_bytes_l impl=impl, pp_layers=[0, 1], num_kv_heads_per_layer=[1, 1], + num_locality_domains=num_locality_domains, get_disagg_role_mapper_kinds=lambda: role_mapper_kinds, get_disagg_role_layouts=lambda: {}, ) +@pytest.mark.parametrize("num_locality_domains", [1, 2]) +def test_v2_builder_emits_mapped_ranges_only_for_localized_pools(num_locality_domains): + """A non-localized pool is fully backed: no ranges, and the wire format is unchanged.""" + from types import SimpleNamespace + + attrs = { + (0, Role.KEY): SimpleNamespace(life_cycle_id=0, pool_index=0, offset=0, size=128), + (0, Role.VALUE): SimpleNamespace(life_cycle_id=0, pool_index=0, offset=128, size=128), + (1, Role.KEY): SimpleNamespace(life_cycle_id=0, pool_index=0, offset=256, size=128), + (1, Role.VALUE): SimpleNamespace(life_cycle_id=0, pool_index=0, offset=384, size=128), + } + ranges = ((0, 1024), (2048, 3072)) + manager = _make_fake_v2_manager( + attrs, + {Role.ALL: MapperKind.NHD}, + slot_bytes_list=(512,), + num_locality_domains=num_locality_domains, + mapped_ranges=ranges, + ) + + pool = build_page_table_from_manager(manager).pool_groups[0].pools[0] + if num_locality_domains == 1: + assert pool.mapped_ranges is None + assert "mapped_ranges" not in pool.to_dict() + else: + assert pool.mapped_ranges == ranges + + def test_v2_builder_stamps_homogeneous_pool_kinds(): """Split KV / INDEX_KEY pools get NHD and REPLICATED views respectively.""" from types import SimpleNamespace