diff --git a/cpp/tensorrt_llm/CMakeLists.txt b/cpp/tensorrt_llm/CMakeLists.txt index 1e13b43ef27a..bbe6127befc4 100644 --- a/cpp/tensorrt_llm/CMakeLists.txt +++ b/cpp/tensorrt_llm/CMakeLists.txt @@ -279,7 +279,7 @@ target_link_libraries(${SHARED_TARGET} PUBLIC ${TRTLLM_LINK_LIBS}) target_link_libraries( ${SHARED_TARGET} PRIVATE - $ + $ $ $ $ @@ -308,6 +308,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 0f033e83ec3b..bd3c59544079 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 19600e0fd52a..3d3429a307e5 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..8bac4c619fe7 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 @@ -66,6 +66,11 @@ Page::~Page() }); } +std::optional Page::localityDomainId() const +{ + return hasValidSlot() ? manager->localityDomainOf(cacheLevel, lifeCycle, slotId()) : std::nullopt; +} + PageStatus Page::status() const noexcept { auto h = holder.lock(); @@ -671,12 +676,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/page.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h index b10075de042f..cadd7906135f 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h @@ -79,6 +79,9 @@ class Page : public Slot, public EnableSharedFromThis PageStatus status() const noexcept; + // Locality domain whose memory backs this page's slot; nullopt off a localized GPU tier. + [[nodiscard]] std::optional localityDomainId() const; + bool scheduledForEviction() const noexcept { return nodeRef.has_value(); 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..1f330120e359 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,7 @@ Slot SlotAllocator::allocate() } else if (mNumActiveSlots < std::min(mCapacity, mTargetCapacity)) { - slot.setSlotId(SlotId{mNumActiveSlots++}); + slot.setSlotId(slotId(mNumActiveSlots++)); } else { @@ -131,7 +168,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 +197,16 @@ 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 (!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 +240,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 +300,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 +326,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 +356,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 +384,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 +404,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 +463,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 +602,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 +623,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 = localityDomainOf(slot.slotId()).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 +713,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 +742,200 @@ 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); + mChunkOwner.assign(mMaxChunks, -1); + + 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); + mChunkOwner[chunk] = static_cast(localityDomainId.value()); + } + while (slotMap.chunks().size() > numChunks) + { + size_t const chunk = slotMap.pop(); + mChunkOwner[chunk] = -1; + mFreeChunks.push(chunk); + } +} + +// 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 +969,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 +984,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 +1030,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 +1041,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 +1061,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..4465dd2c1bba 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 @@ -129,14 +131,59 @@ struct Slot std::optional mSlotId; }; +// --------------------------------------------------------------------------- +// 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 +194,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 +250,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 +335,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 +367,7 @@ class GpuSlotPool : public SlotPoolBase private: VirtMem mVirtMem; + TypedVec const* mSlotMaps; }; // --------------------------------------------------------------------------- @@ -354,28 +443,57 @@ 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); + + // Locality domain whose memory backs a slot id this group issued; nullopt when not localized. + [[nodiscard]] virtual std::optional localityDomainOf(SlotId /*slot*/) const noexcept + { + return 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 +511,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 +521,68 @@ 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(); + } + + [[nodiscard]] std::optional localityDomainOf(SlotId slot) const noexcept override + { + if (mSlotMaps.empty()) + { + return std::nullopt; + } + int8_t const owner = mChunkOwner[slotCountToSizeT(slot.value() / mChunkSize)]; + return owner < 0 ? std::nullopt : std::optional{LocalityDomainId{owner}}; + } + +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; + // Owning locality domain of each chunk of the slot-id space, kept in step with mSlotMaps. + std::vector mChunkOwner; }; //! Host pools are committed progressively on a background thread, so a slot may @@ -439,8 +604,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 +635,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 +659,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)->slotIndexUpperBound(); + } + + std::vector> slotRanges(PoolGroupIndex pgIdx) const { - return mPoolGroups.at(pgIdx)->numSlots(); + return mPoolGroups.at(pgIdx)->slotRanges(); } TypedVec slotSize(PoolGroupIndex pgIdx) const @@ -506,6 +689,11 @@ class CacheLevelStorage return *mPoolGroups.at(pgIdx); } + PoolGroupBase const& poolGroup(PoolGroupIndex pgIdx) const + { + return *mPoolGroups.at(pgIdx); + } + MemAddress getBaseAddress(PoolGroupIndex pgIdx, PoolIndex poolIdx, SlotId slotId) const { return std::get(mPoolGroups.at(pgIdx)->slotAddress(slotId).at(poolIdx)); @@ -520,13 +708,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 +806,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..fb1efda6996e 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 @@ -746,6 +804,11 @@ void StorageManager::releaseSlot(LifeCycleId lc, CacheLevel level, Slot slot) mLevels.at(level).storage->release(pg, std::move(slot)); } +std::optional StorageManager::localityDomainOf(CacheLevel level, LifeCycleId lc, SlotId slot) const +{ + return mLevels.at(level).storage->poolGroup(getPoolGroupIndex(level, lc)).localityDomainOf(slot); +} + // --------------------------------------------------------------------------- // isEvictable // --------------------------------------------------------------------------- @@ -878,14 +941,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 +963,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 +1023,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 +1085,8 @@ void StorageManager::_prepareFreeSlots(TypedVec const domainFilter = localityDomainFilter(lvlId, localityDomainId); auto countPages = [&](PoolGroupIndex pgIdx, PagesByLifeCycle const& pagesByLifeCycle) { @@ -1036,8 +1103,8 @@ void StorageManager::_prepareFreeSlots(TypedVec oldFree ? required - oldFree : 0; numToEvict.at(pgIdx) = std::min(shortage, evictableCount); @@ -1074,7 +1141,7 @@ void StorageManager::_prepareFreeSlots(TypedVec>> acceptedPages(numLifeCycles()); @@ -1116,7 +1183,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 +1242,7 @@ void StorageManager::_prepareFreeSlots(TypedVecstatus() == PageStatus::DROPPABLE) @@ -1244,7 +1312,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 +1343,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 +1444,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 +1460,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 +1472,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 +1547,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 +1566,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 +1658,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 +1705,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 +1738,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 +1824,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 +1834,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 +1857,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 +1922,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 +1974,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 +1997,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 +2028,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 +2042,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 +2059,7 @@ void StorageManager::shrinkPoolGroup( // Finalize shrink and resize pools. allocator.finishShrink(); - pg.resizePools(newNumSlots); + pg.resizePools(newNumSlots, localityDomainId); } // --------------------------------------------------------------------------- @@ -1952,10 +2071,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 +2104,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 +2461,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..dc8c06536a8f 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,20 +143,26 @@ 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); + // Locality domain whose memory backs a slot of this life cycle at this level; nullopt when not localized. + [[nodiscard]] std::optional localityDomainOf(CacheLevel level, LifeCycleId lc, SlotId slot) const; // ---- Eviction ---------------------------------------------------------- @@ -175,27 +181,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 +214,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 +245,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 +338,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 +423,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 +509,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/common/attentionOp.cpp b/cpp/tensorrt_llm/common/attentionOp.cpp index f93d1c637ad3..0c348a26b1bb 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/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index c503e0250648..9bcc22bbb16f 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); @@ -2114,6 +2143,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()) @@ -2546,6 +2581,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) @@ -2900,6 +2954,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/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/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/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/tensorrt_llm/thop/attentionOp.cpp b/cpp/tensorrt_llm/thop/attentionOp.cpp index 57e644a1a639..689f3d87a3ff 100644 --- a/cpp/tensorrt_llm/thop/attentionOp.cpp +++ b/cpp/tensorrt_llm/thop/attentionOp.cpp @@ -1144,9 +1144,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; @@ -1246,10 +1247,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 @@ -1492,7 +1494,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 75c228ebb970..e01070e2ca8b 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; @@ -197,13 +229,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/cpp/tests/unit_tests/batch_manager/CMakeLists.txt b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt index 0f2879b1dd77..b09188a308b0 100644 --- a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt +++ b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt @@ -264,6 +264,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..b60af0c0ed39 --- /dev/null +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2LocalityTest.cpp @@ -0,0 +1,559 @@ +/* + * 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, MappedAllocatorsIssueDisjointGlobalIds) +{ + 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_LT(slot.slotId(), SlotId{4}); + ids.insert(slot.slotId()); + } + for (auto const& slot : slots1) + { + EXPECT_GE(slot.slotId(), SlotId{4}); + ids.insert(slot.slotId()); + } + EXPECT_EQ(ids.size(), 8U); + EXPECT_EQ(alloc0.numFreeSlots(), 0); + + // A slot id from the other domain's chunk is foreign to this allocator; the real slot stays + // with alloc1 and is released below. + { + Slot foreign; + foreign.setSlotId(slots1[0].slotId()); + 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's locality domain comes from the chunk its id falls in, so it stays right after one domain +// returns chunks and the other domain takes them over. +TEST_F(KvCacheManagerV2MockLocalityTest, SlotIdDeterminesTheLocalityDomain) +{ + auto manager = std::make_shared(makeLocalizedConfig()); + auto& storage = manager->storage(); + PoolGroupIndex const pg{0}; + LifeCycleId const lc{0}; + SlotCount const before = storage.numSlots(pg); + storage.shrinkPoolGroup(kHotLevel, pg, SlotCount{1}, {}, LocalityDomainId{0}); + storage.expandPoolGroup(kHotLevel, pg, 2 * before, LocalityDomainId{1}); + storage.expandPoolGroup(kHotLevel, pg, before, LocalityDomainId{0}); + + auto const freeBefore = freeSlotsPerDomain(storage); + for (LocalityDomainId domain{0}; domain < storage.numLocalityDomains(); ++domain) + { + auto slots = storage.newSlotsForPoolGroup(kHotLevel, pg, freeBefore[domain.value()], {}, {}, domain); + for (auto& slot : slots) + { + ASSERT_EQ(storage.localityDomainOf(kHotLevel, lc, slot.slotId()), std::optional{domain}); + storage.releaseSlot(lc, kHotLevel, std::move(slot)); + } + } + EXPECT_EQ(freeSlotsPerDomain(storage), freeBefore); +} + +// 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/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. diff --git a/docs/source/features/locality-domains.md b/docs/source/features/locality-domains.md new file mode 100644 index 000000000000..ed34e7e2a963 --- /dev/null +++ b/docs/source/features/locality-domains.md @@ -0,0 +1,18 @@ +# Locality Domains + +A locality domain is one half of a Rubin GPU: a set of SMs and the HBM attached to them. With locality +domains enabled, each request's KV cache lives in one domain and its attention runs on that domain's SMs, +so KV cache reads do not cross the die-to-die link. This is a prototype and is off by default. + +Consider enabling it for dense attention variants such as GQA and MLA at high batch sizes or high +concurrency, where decode attention is bandwidth-bound. Leave it off for low concurrency or small decode +batches, where splitting each step across the two domains costs more than it saves. Sparse attention models +are not supported. + +```yaml +enable_locality_domains: true +kv_cache_config: + use_kv_cache_manager_v2: true +``` + +Prefix reuse happens only within a domain. diff --git a/docs/source/index.rst b/docs/source/index.rst index 9920583eb754..ae5823f10eab 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -71,6 +71,7 @@ Welcome to TensorRT LLM's Documentation! features/subagent-routing.md features/embeddings.md features/kvcache.md + features/locality-domains.md features/kv-cache-compression.md features/long-sequence.md features/lora.md diff --git a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py index acc37b4b91be..b28c304b46ca 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 5c2bbb8e12dd..06edf0d4fb16 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, is_sm_100f, is_sm_120f, maybe_pin_memory, prefer_pinned) from tensorrt_llm.bindings.internal import thop @@ -41,6 +44,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) @@ -61,6 +65,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, @@ -316,6 +326,37 @@ 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 + # Per-domain (_cute_dsl_mla_split_plan, _cute_dsl_mla_split_plan_key). + locality_domain_cute_dsl_mla_split_plans: Optional[List[Tuple[ + Optional[torch.Tensor], Optional[tuple]]]] = 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: """ @@ -482,6 +523,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", @@ -695,6 +738,33 @@ 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) + ] + if (self.locality_domain_cute_dsl_mla_split_plans is None + or len(self.locality_domain_cute_dsl_mla_split_plans) + != num_locality_domains): + self.locality_domain_cute_dsl_mla_split_plans = [ + (None, None) + ] * num_locality_domains + def update_for_spec_dec(self) -> None: if self.enable_flash_mla: self._flash_mla_metadata_valid = False @@ -709,6 +779,95 @@ def _invalidate_mla_scheduler_buffers(self) -> None: self._mla_scheduler_buffers_valid = False self._mla_ctx_cu_seqlens_valid = False self._cute_dsl_mla_split_plan_key = None + plans = self.locality_domain_cute_dsl_mla_split_plans + if plans is not None: + plans[:] = [(plan, None) for plan, _ in plans] + + 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() @@ -820,6 +979,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, *, @@ -1125,6 +1329,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 @@ -1140,6 +1402,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, @@ -1528,6 +1796,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 @@ -1624,10 +1907,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 @@ -1702,6 +1987,646 @@ def is_sm_version_trtllm_gen_kernel(self, sm): return is_sm_100f(sm) +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] + # CuTe DSL MLA caches its split-KV plan on the metadata and rewrites it when + # the key changes; a shared plan would be rewritten by one domain's stream + # while the other domain's kernels still read it. + plans = metadata.locality_domain_cute_dsl_mla_split_plans + has_plan = plans is not None and locality_domain_id < len(plans) + if has_plan: + (overrides["_cute_dsl_mla_split_plan"], + overrides["_cute_dsl_mla_split_plan_key"]) = plans[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: + if has_plan: + plans[locality_domain_id] = (metadata._cute_dsl_mla_split_plan, + metadata._cute_dsl_mla_split_plan_key) + 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 @@ -2507,6 +3432,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_() @@ -2562,29 +3492,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 @@ -2598,6 +3531,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 @@ -2990,6 +4114,28 @@ def _mla_rope_generation_impl( q_rope_applied, ) + # 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/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/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 151c4d84ddcb..4830e6bb561b 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2081,6 +2081,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 self._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 @@ -2115,6 +2120,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, @@ -2439,6 +2445,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, @@ -3297,6 +3304,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, @@ -3453,6 +3461,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/engine/runners/decoder/dummy_requests.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/dummy_requests.py index 3837bc138d6f..36c10f0412c6 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, @@ -343,45 +344,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 @@ -391,10 +427,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 ) @@ -409,6 +444,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( ResourceManagerType.KV_CACHE_MANAGER ) @@ -422,7 +479,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, @@ -432,13 +490,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, ) @@ -457,9 +517,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, @@ -467,13 +527,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 a85547af6ea8..f06ca874db2a 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py @@ -59,6 +59,7 @@ from .config import DecoderRunnerConfig from .dummy_requests import DummyRequestFactory from .inputs import DecoderInputPreparer +from .locality_domain import locality_domain_forkjoin_enabled from .metadata import DecoderMetadata from .warmup import WarmupAgreement @@ -454,6 +455,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 2ec483ecba15..004204b1c373 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py @@ -6,6 +6,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 @@ -66,6 +67,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 @@ -286,6 +292,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 @@ -856,6 +864,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, @@ -1038,6 +1047,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) @@ -1238,6 +1257,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: @@ -2280,6 +2306,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 @@ -2301,6 +2331,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/pyexecutor/executor_request_queue.py b/tensorrt_llm/_torch/pyexecutor/executor_request_queue.py index 8cc8d971ca4f..3eef0babd471 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 8984856fd55e..13db6c1fad63 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 ( @@ -1755,6 +1756,42 @@ def copy_block_offsets_to( ) +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 @@ -1827,6 +1864,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, @@ -2232,7 +2270,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. @@ -2393,6 +2433,17 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: ) 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 + # Live KV caches per locality domain, the load that prefix-aware placement balances against. + self._locality_domain_load = [0] * self.num_locality_domains self.can_evict = len(config.cache_tiers) > 1 if self.event_manager is not None: self.event_manager.set_layer_group_window_sizes( @@ -2677,6 +2728,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") @@ -4538,6 +4599,93 @@ 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, + input_tokens: Sequence[int] | None = None, + lora_task_id: int | None = None, + cache_salt: str | None = None, + ) -> int | None: + """Locality domain placement for a newly-created request. + + Prefixes are reused only within a domain, so a request goes to the domain + holding the longest cached prefix of its prompt when that match covers at + least one block and the domain is not already carrying clearly more live + requests than the others. Otherwise placement is round-robin over arrival + order rather than by request id: ids that advance in even steps (a parent + and its child per request, or round-robin disagg ids) would pin every + request to one domain. + """ + if self.num_locality_domains <= 1: + return None + if input_tokens is not None and self.enable_block_reuse: + matches = self._prefix_match_per_domain(input_tokens, lora_task_id, cache_salt) + best = max(range(self.num_locality_domains), key=matches.__getitem__) + load = self._locality_domain_load + if matches[best] >= self.tokens_per_block and load[best] - min(load) <= max( + 8, sum(load) // 4 + ): + return best + 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) + # The target manager places the request, and a separate draft manager can + # resolve a different domain count on its own slot sizes and quota. The + # draft sub-forward then runs unpartitioned, so the draft only needs a + # valid, stable id of its own. + # TODO: resolve one domain count for the target and draft managers at + # initialization, then drop the is_draft cases here. + if self.num_locality_domains <= 1: + assert locality_domain_id is None or locality_domain_id == 0 or self.is_draft, ( + "Non-localized requests must not use a nonzero locality_domain_id" + ) + return None + if locality_domain_id is None and self.is_draft: + locality_domain_id = req.py_request_id % self.num_locality_domains + 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: @@ -5049,7 +5197,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), ) @@ -5113,6 +5265,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, @@ -5124,9 +5277,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: @@ -5975,6 +6133,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 @@ -6806,6 +6965,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. @@ -6827,6 +6987,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 ) @@ -6881,6 +7045,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 = LlmRequest( request_id=req_id, max_new_tokens=1, @@ -6888,6 +7057,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 @@ -6901,6 +7071,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. @@ -6923,16 +7094,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) @@ -7131,6 +7307,8 @@ def free_resources(self, request: LlmRequest, pin_on_release: bool = False): kv_cache.discard_pending_stats() if self._sharing is not None: self._sharing._on_free(request.py_request_id, kv_cache) + if self.num_locality_domains > 1: + self._locality_domain_load[kv_cache.locality_domain_id] -= 1 kv_cache.close() self.impl.clear_stats_excluded(request.py_request_id) if request.py_request_id in self._early_freed_index_requests: @@ -7377,9 +7555,62 @@ 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(): + if self.num_locality_domains > 1: + pool_id = self.layer_to_pool_mapping_dict[layer_offset] + if pool_id in pool_handled: + continue + # 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) + continue buffer = self.get_buffers(layer_id) if buffer is None: continue @@ -8130,10 +8361,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 " @@ -8165,7 +8409,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, @@ -8174,6 +8420,8 @@ def _create_kv_cache( **priority_kwargs, ) self.kv_cache_map[request_id] = kv_cache + if locality_domain_id is not None: + self._locality_domain_load[locality_domain_id] += 1 if enable_request_stats and not self.enable_stats: self._request_stats_enabled_ids.add(request_id) if is_dummy: @@ -8202,11 +8450,30 @@ def probe_prefix_match_length(self, input_tokens, lora_task_id=None, cache_salt= return 0 if len(input_tokens) == 0: return 0 + # The router has no request yet, so no domain: blocks are keyed per + # domain and placement sends the request to the best-matching one. + return max(self._prefix_match_per_domain(input_tokens, lora_task_id, cache_salt)) + + def _prefix_match_per_domain( + self, input_tokens, lora_task_id=None, cache_salt=None + ) -> list[int]: + """Cached prefix length of *input_tokens* in each locality domain (one entry if not localized).""" + if not self.enable_block_reuse or len(input_tokens) == 0: + return [0] * max(1, self.num_locality_domains) 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 + ) + ] + return [ + 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: """Key of the first block whose pages *req* would commit, or None. @@ -8248,7 +8515,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: @@ -8268,7 +8539,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/mamba_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py index 5019ed8b8a33..0bcb070a5f40 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py @@ -882,6 +882,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, @@ -892,6 +893,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 @@ -930,6 +932,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, @@ -942,6 +945,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 a4e42209a2de..0fb069c46053 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -718,6 +718,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) @@ -807,6 +808,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 @@ -1250,6 +1254,7 @@ def create_child_request(self, child_id: int) -> None: is not None else None), py_multimodal_data=deepcopy(self.py_multimodal_data), encoder_input_features=self.py_encoder_input_features, + locality_domain_id=self.py_locality_domain_id, ) child.py_bad_words = deepcopy(self.py_bad_words) child.py_no_repeat_ngram_size = self.py_no_repeat_ngram_size @@ -1568,7 +1573,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 @@ -1726,6 +1732,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 87f83a022cb5..479d5303070f 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -7267,6 +7267,18 @@ 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: + request = req_item.request + lora_config = getattr(request, "lora_config", None) + req_item.locality_domain_id = ( + self.kv_cache_manager.pick_locality_domain( + getattr(request, "input_token_ids", None), + lora_config.task_id + if lora_config is not None else None, + getattr(request, "cache_salt", None))) + # 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 aa8f5636bd6c..6224c68ef841 100644 --- a/tensorrt_llm/_torch/pyexecutor/request_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/request_utils.py @@ -268,7 +268,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: @@ -312,6 +316,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 4cf174f17fd0..b025cb23a2ef 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -630,7 +630,12 @@ def preempt_for_pages(req: LlmRequest) -> bool: if other.py_request_id in demoted_ids: protected.add(other.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 --- @@ -1804,6 +1809,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. @@ -1833,6 +1839,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/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 97b67a555523..2f5b4ca6adf7 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -3250,19 +3250,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 @@ -3276,13 +3297,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/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index e418a261853e..c5ad8cc1468a 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: ... @@ -670,6 +674,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: ... @@ -726,6 +732,7 @@ class PoolDesc: pool_index: PoolIndex base_address: MemAddress slot_bytes: int + mapped_ranges: list[tuple[int, int]] class BufferId(NamedTuple): layer_id: LayerId @@ -915,6 +922,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/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/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 4727ba8b8e40..43ecc86d376c 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -133,6 +133,8 @@ l0_b200: - unittest/_torch/kv_cache_compression - unittest/disaggregated/test_deepseek_v4_kv_transfer.py - unittest/disaggregated/test_fp4_mla_kv_transfer.py + # Locality domains via mock localization: per-domain KV pools on one GPU. + - unittest/disaggregated/test_cache_transceiver_single_process.py::test_cache_transceiver_locality_domains - unittest/_torch/attention/sparse/deepseek_v4/test_prims_ts_sparse_mla_fmha.py TIMEOUT (60) - unittest/_torch/modeling/test_modeling_deepseekv4.py - unittest/_torch/modeling/test_glm5_next_contracts.py diff --git a/tests/unittest/_torch/attention/test_attention_mla.py b/tests/unittest/_torch/attention/test_attention_mla.py index 9c3c83e284d7..79b8d8d0be8c 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_fmha_interface.py b/tests/unittest/_torch/attention/test_fmha_interface.py index 6c4f74997307..3258c2fbd1a1 100644 --- a/tests/unittest/_torch/attention/test_fmha_interface.py +++ b/tests/unittest/_torch/attention/test_fmha_interface.py @@ -14,6 +14,7 @@ # limitations under the License. from contextlib import nullcontext +from types import MethodType from typing import cast from unittest.mock import Mock, patch @@ -81,6 +82,7 @@ def test_mla_forward_clears_fused_qkv_before_fmha_selection( attn.kv_scale_quant_orig = None attn.get_local_layer_idx.return_value = 0 attn._ensure_rope_table_size = Mock() + attn._attention_partitions = MethodType(TrtllmAttention._attention_partitions, attn) fmha = Mock() attn._fmha_manager = Mock() attn._fmha_manager.select.return_value = fmha @@ -94,6 +96,7 @@ def test_mla_forward_clears_fused_qkv_before_fmha_selection( metadata.spec_decoding_bl_tree_mask = None metadata.max_context_q_len_override = None metadata.kv_cache_manager = None + metadata.locality_domain_enabled = False lengths = torch.ones(1, dtype=torch.int32) metadata.kv_lens_cuda_runtime = lengths metadata.kv_lens_runtime = lengths 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..bbfaf48e8f6a --- /dev/null +++ b/tests/unittest/_torch/attention/test_locality_domain_attention_gpu.py @@ -0,0 +1,342 @@ +# 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 _cached_lengths(cached, batch: int) -> list[int]: + return list(cached) if isinstance(cached, (list, tuple)) else [cached] * batch + + +def _manager( + *, mla: bool, batch: int, cached, num_kv_heads: int, head_dim: int, localized: bool +): + tokens_per_block = 32 if mla else 64 + lengths = _cached_lengths(cached, batch) + max_seq_len = max(lengths) + 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)), + [length + 1 for length in lengths], + 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, localized: bool): + lengths = _cached_lengths(cached, batch) + kwargs = dict( + num_contexts=0, + kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=lengths), + 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=[length + 1 for length in lengths], + 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, localized: bool, num_heads: int = 64) -> torch.Tensor: + gen = torch.Generator(device="cuda").manual_seed(0) + latent = [ + torch.randn(length, _MLA_HEAD_DIM, device="cuda", generator=gen) * 0.5 + for length in _cached_lengths(cached, 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) + + +def test_cute_dsl_mla_decode_with_unequal_domain_lengths(monkeypatch): + # CuTe DSL MLA caches its split-KV plan on the metadata. Domain 0 holds short + # contexts and domain 1 long ones, so their plans differ; with one shared plan + # the domain-1 writer would race domain 0's kernels on the other stream. + from tensorrt_llm._torch.attention.backends.fmha.cute_dsl_mla import CuteDslMlaFmha + + monkeypatch.setenv("TLLM_FMHA_LIBS", "cute_dsl_mla,fallback") + plans = [] + split_kv_plan = CuteDslMlaFmha._split_kv_plan + + def recording_split_kv_plan(*args, **kwargs): + plan = split_kv_plan(*args, **kwargs) + plans.append(None if plan is None else plan.data_ptr()) + return plan + + monkeypatch.setattr(CuteDslMlaFmha, "_split_kv_plan", staticmethod(recording_split_kv_plan)) + batch = 32 + cached = [1000] * (batch // 2) + [16384] * (batch // 2) + off = _mla_output(batch=batch, cached=cached, localized=False, num_heads=96) + assert len(plans) == 1 and plans[0] is not None + plans.clear() + on = _mla_output(batch=batch, cached=cached, localized=True, num_heads=96) + assert len(plans) == 2 and None not in plans + assert plans[0] != plans[1] + 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..c67f669e85e2 --- /dev/null +++ b/tests/unittest/_torch/attention/test_locality_domain_partitions.py @@ -0,0 +1,476 @@ +# 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_cute_dsl_mla_split_plans = None + self._cute_dsl_mla_split_plan = None + self._cute_dsl_mla_split_plan_key = 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 test_scoped_metadata_gives_each_domain_its_own_cute_dsl_mla_split_plan(): + metadata = _Metadata(num_ctx=0, num_gen=4, gen_tokens=(1,) * 4) + full_plan = torch.zeros(4, dtype=torch.int32) + metadata._cute_dsl_mla_split_plan = full_plan + metadata._cute_dsl_mla_split_plan_key = ("full",) + metadata.locality_domain_cute_dsl_mla_split_plans = [(None, None), (None, None)] + slices = [ + _slice(ctx_req=(0, 0), gen_req=(0, 2), gen_tok=(0, 2), n=2), + _slice(ctx_req=(0, 0), gen_req=(2, 4), gen_tok=(2, 4), n=2), + ] + written = [] + for domain, ld_slice in enumerate(slices): + with trtllm_backend._locality_domain_metadata(metadata, ld_slice, domain): + # Each domain starts from its own (empty) plan, never another domain's. + assert metadata._cute_dsl_mla_split_plan is None + assert metadata._cute_dsl_mla_split_plan_key is None + plan = torch.full((2,), domain, dtype=torch.int32) + metadata._cute_dsl_mla_split_plan = plan + metadata._cute_dsl_mla_split_plan_key = (domain,) + written.append(plan) + assert metadata._cute_dsl_mla_split_plan is full_plan + assert metadata._cute_dsl_mla_split_plan_key == ("full",) + plans = metadata.locality_domain_cute_dsl_mla_split_plans + assert [plan for plan, _ in plans] == written + assert [key for _, key in plans] == [(0,), (1,)] + # A later layer of the same step reuses its domain's plan. + with trtllm_backend._locality_domain_metadata(metadata, slices[1], 1): + assert metadata._cute_dsl_mla_split_plan is written[1] + assert metadata._cute_dsl_mla_split_plan_key == (1,) + + metadata.locality_domain_enabled = False + trtllm_backend.TrtllmAttentionMetadata._invalidate_mla_scheduler_buffers(metadata) + assert metadata._cute_dsl_mla_split_plan_key is None + assert [plan for plan, _ in plans] == written + assert [key for _, key in plans] == [None, None] + + +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/engine/runners/decoder/test_dummy_requests.py b/tests/unittest/_torch/executor/engine/runners/decoder/test_dummy_requests.py index 6c0049d8376c..d4ad87718439 100644 --- a/tests/unittest/_torch/executor/engine/runners/decoder/test_dummy_requests.py +++ b/tests/unittest/_torch/executor/engine/runners/decoder/test_dummy_requests.py @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from unittest.mock import Mock, PropertyMock, patch import pytest import torch @@ -8,6 +9,7 @@ from _torch.executor.engine_test_utils import DummyModel, kv_cache_manager from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.dummy_requests import DummyRequestFactory +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType from tensorrt_llm.llmapi.llm_args import ( DecodingBaseConfig, @@ -60,3 +62,61 @@ def test_warmup_builders_resynchronize_stale_draft_buffers( len(request.py_draft_tokens) == expected_buffer_width for request in batch.generation_requests ) + + +def _warmup_request_split(num_locality_domains: int): + model = DummyModel(torch.half) + requests = DummyRequestFactory( + config=decoder_config(), + model=model, + get_runtime_tokens_per_gen_step=lambda draft_len: draft_len + 1, + release_request_state=None, + ) + with kv_cache_manager(model=model) as manager: + manager.get_num_free_blocks = Mock(return_value=8) + manager.max_seq_len = 2048 + 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) + manager.add_dummy_requests = Mock( + side_effect=[ + [_dummy_request() for _ in range(4 - num_long)], + [_dummy_request() for _ in range(num_long)], + ] + ) + resource_manager = ResourceManager({ResourceManagerType.KV_CACHE_MANAGER: manager}) + with patch.object( + KVCacheManagerV2, + "num_locality_domains", + new_callable=PropertyMock, + return_value=num_locality_domains, + ): + requests.create_cuda_graph_warmup_request( + resource_manager, batch_size=4, draft_len=0, max_seq_len=512 + ) + return manager.add_dummy_requests.call_args_list + + +def test_cuda_graph_warmup_request_split(): + short_call, long_call = _warmup_request_split(1) + assert short_call.args[0] == [0, 1, 2] + assert "locality_domain_ids" not in short_call.kwargs + assert long_call.kwargs["request_ids"] == [3] + assert long_call.kwargs["token_nums"] == [100] + assert "locality_domain_ids" not in long_call.kwargs + + +def test_cuda_graph_warmup_request_split_per_locality_domain(): + # One longest request per domain, short requests dealt from domain 0. + short_call, long_call = _warmup_request_split(2) + assert short_call.args[0] == [0, 1] + assert short_call.kwargs["locality_domain_ids"] == [0, 1] + assert long_call.kwargs["request_ids"] == [2, 3] + assert long_call.kwargs["token_nums"] == [100, 100] + assert long_call.kwargs["locality_domain_ids"] == [0, 1] 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 2b1eb05dd600..249c71520f6d 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 @@ -82,7 +82,9 @@ def _make_creator( c._mapping = Mock() c._draft_config = None c._model_engine = Mock(cuda_graph_batch_sizes=[], cuda_graph_padding_enabled=False) - 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() @@ -379,6 +381,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) @@ -1180,7 +1185,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 ef1f2d68b17b..013fc858baeb 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 @@ -1743,6 +1743,7 @@ def _one_model_draft_creator(target_pool_ratio, mapping, draft_config) -> KvCach creator._draft_config = draft_config creator._speculative_config = Mock() creator._disable_overlap_scheduler = False + creator._llm_args = SimpleNamespace(enable_locality_domains=False) return creator 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 ae062f0983c5..a6469b1ca2f9 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 @@ -2275,6 +2275,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_context_only_warmup.py b/tests/unittest/_torch/executor/test_context_only_warmup.py index 83fd66d08c40..eaee6bf8c850 100644 --- a/tests/unittest/_torch/executor/test_context_only_warmup.py +++ b/tests/unittest/_torch/executor/test_context_only_warmup.py @@ -24,6 +24,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_disagg_receive_ordering.py b/tests/unittest/_torch/executor/test_disagg_receive_ordering.py index 89d9e727ff11..95cf10cfce38 100644 --- a/tests/unittest/_torch/executor/test_disagg_receive_ordering.py +++ b/tests/unittest/_torch/executor/test_disagg_receive_ordering.py @@ -62,6 +62,7 @@ def _manager() -> Mock: manager.kv_connector_manager = None manager.enable_block_reuse = False manager._has_cp_helix = False + manager.num_locality_domains = 1 manager.num_extra_kv_tokens = 0 manager.kv_cache_map = {} manager.prepare_context_cache.return_value = 0 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_locality_domain_request.py b/tests/unittest/_torch/executor/test_locality_domain_request.py new file mode 100644 index 000000000000..cd415a844f6d --- /dev/null +++ b/tests/unittest/_torch/executor/test_locality_domain_request.py @@ -0,0 +1,136 @@ +# 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_round_robins_over_arrival_order(): + manager = SimpleNamespace(num_locality_domains=2, _next_locality_domain=0) + picks = [KVCacheManagerV2.pick_locality_domain(manager) for _ in range(4)] + assert picks == [0, 1, 0, 1] + + +def _placer(matches, load=(0, 0)): + return SimpleNamespace( + num_locality_domains=2, + enable_block_reuse=True, + tokens_per_block=32, + _next_locality_domain=0, + _locality_domain_load=list(load), + _prefix_match_per_domain=lambda *args: list(matches), + ) + + +def test_placement_follows_the_domain_holding_the_prompt_prefix(): + manager = _placer(matches=(0, 96)) + tokens = list(range(128)) + assert KVCacheManagerV2.pick_locality_domain(manager, tokens) == 1 + assert KVCacheManagerV2.pick_locality_domain(manager, tokens) == 1 + + +def test_placement_round_robins_without_a_full_block_match(): + manager = _placer(matches=(0, 31)) + picks = [KVCacheManagerV2.pick_locality_domain(manager, list(range(64))) for _ in range(4)] + assert picks == [0, 1, 0, 1] + + +def test_placement_does_not_pile_a_hot_prefix_onto_a_loaded_domain(): + # Domain 1 already holds 30 more live requests than domain 0, past max(8, total // 4). + manager = _placer(matches=(0, 96), load=(10, 40)) + assert KVCacheManagerV2.pick_locality_domain(manager, list(range(128))) == 0 + + +def _manager(num_locality_domains, is_draft): + manager = object.__new__(KVCacheManagerV2) + manager.impl = SimpleNamespace(num_locality_domains=num_locality_domains) + manager.is_draft = is_draft + manager.enable_block_reuse = False + manager._next_locality_domain = 0 + return manager + + +@pytest.mark.parametrize("target_domains, draft_domains", [(2, 1), (1, 2)]) +def test_draft_manager_with_a_different_domain_count_gets_a_valid_domain( + target_domains, draft_domains +): + target = _manager(target_domains, is_draft=False) + draft = _manager(draft_domains, is_draft=True) + for request_id in range(4): + request = _request(request_id, target.pick_locality_domain()) + assert target._request_locality_domain_id(request) == request.py_locality_domain_id + draft_domain = draft._request_locality_domain_id(request) + if draft_domains == 1: + assert draft_domain is None + else: + assert 0 <= draft_domain < draft_domains + assert draft._request_locality_domain_id(request) == draft_domain diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index cc65c61de6ca..95cdcfb0cd80 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -3120,6 +3120,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) @@ -3206,6 +3207,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 304eb96f50c1..642f344f0eb8 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/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py index 8cebeff8c38e..af420a2b4e7d 100644 --- a/tests/unittest/_torch/speculative/test_eagle3.py +++ b/tests/unittest/_torch/speculative/test_eagle3.py @@ -80,7 +80,7 @@ def on_update_kv_lens(self): metadata = _Metadata() @contextmanager - def draft_context(attn_metadata, manager): + def draft_context(attn_metadata, manager, enable_draft_locality_domain_forkjoin=False): assert attn_metadata is metadata assert manager is draft_manager events.append("switch") 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 b2c33ae83ce3..da92ffc34efa 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() + ) 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 ec6ef42ad0f7..e21539c9156e 100644 --- a/tests/unittest/disaggregated/test_extractor.py +++ b/tests/unittest/disaggregated/test_extractor.py @@ -253,7 +253,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): @@ -291,7 +299,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( @@ -314,11 +324,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 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..ace4a71853ee --- /dev/null +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_localization_wrapper.py @@ -0,0 +1,185 @@ +# 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 = {} + wrapper.tokens_per_block = impl.tokens_per_block + wrapper._next_locality_domain = 0 + wrapper._locality_domain_load = [0, 0] + 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_placement_sends_a_request_to_the_domain_holding_its_prefix(self) -> None: + # Round-robin would start at domain 0, so place another request first to move it to 1. + self.assertEqual(self.wrapper.pick_locality_domain(), 0) + self.assertEqual(self.wrapper.pick_locality_domain(self.tokens), 0) + self.assertEqual(self.wrapper.pick_locality_domain([10**6 + t for t in self.tokens]), 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") diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index ede5add24e59..46fd0a40a7b5 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")