[TRTLLM-17239][feat] Locality domain support for the C++ KV cache manager v2 and two-stream attention

cpp/tensorrt_llm/CMakeLists.txt

@@ -279,7 +279,7 @@ target_link_libraries(${SHARED_TARGET} PUBLIC ${TRTLLM_LINK_LIBS})
 target_link_libraries(
   ${SHARED_TARGET}
   PRIVATE
-    $<LINK_LIBRARY:WHOLE_ARCHIVE,tllm_common_error,tllm_common_logging,tllm_common_cuda,tllm_executor_types,tllm_executor_serialization,tllm_runtime_memory,tllm_runtime_buffers,tllm_runtime_distributed,tllm_runtime_topology>
+    $<LINK_LIBRARY:WHOLE_ARCHIVE,tllm_common_error,tllm_common_logging,tllm_common_cuda,tllm_executor_types,tllm_executor_serialization,tllm_runtime_memory,tllm_runtime_buffers,tllm_runtime_distributed,tllm_runtime_topology,tllm_runtime_locality_domain>
     $<LINK_LIBRARY:WHOLE_ARCHIVE,tllm_common_environment,tllm_executor_values,tllm_executor_transfer_agent,tllm_kv_cache_host_memory,tllm_kv_cache,tllm_kv_cache_legacy,tllm_batch_scheduler,tllm_kernel_kv_cache_copy>
     $<LINK_LIBRARY:WHOLE_ARCHIVE,${BATCH_MANAGER_TARGET}>
     $<LINK_LIBRARY:WHOLE_ARCHIVE,${EXECUTOR_TARGET}>
@@ -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

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

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

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<uint8_t const*>(&length), sizeof(length));
         emit(reinterpret_cast<uint8_t const*>(scope.salt->data()), scope.salt->size());
     }
+    if (scope.localityDomainId.has_value())
+    {
+        std::uint64_t const value = static_cast<std::uint64_t>(scope.localityDomainId->value());
+        emit(reinterpret_cast<uint8_t const*>(&value), sizeof(value));
+    }
 }
 
 // ---------------------------------------------------------------------------

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<LoraTaskIdType> loraId;
     std::optional<std::string> 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> 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;
     }
 };
 

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/coldPageCodec.h

@@ -25,6 +25,8 @@
 #include <cstdint>
 #include <memory>
 #include <type_traits>
+#include <utility>
+#include <vector>
 
 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<std::pair<size_t, size_t>> mappedRanges;
 };
 
 //! Describes one hot pool group and the lifecycles that share its physical layout.

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.h

@@ -114,6 +114,9 @@ inline constexpr FileDescriptor kBadFileDescriptor = -1;
 using PageIndex = StrongIndex<int, struct PageIndexTag, -1>;
 inline constexpr PageIndex kBadPageIndex{-1};
 
+// Index of a GPU locality domain.
+using LocalityDomainId = StrongIndex<int, struct LocalityDomainIdTag>;
+
 // Eviction priority (0 = highest priority to evict, 100 = lowest).
 using Priority = int;
 inline constexpr Priority kPriorityMin = 0;

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
     {