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

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/CMakeLists.txt

@@ -50,6 +50,7 @@ set(KV_CACHE_MANAGER_V2_SRCS
     kv_cache_manager_v2/utils/staticHostMem.cpp
     kv_cache_manager_v2/utils/poison.cpp
     kv_cache_manager_v2/cudaVirtMem.cpp
+    kv_cache_manager_v2/localityDomain.cpp
     kv_cache_manager_v2/storage/config.cpp
     kv_cache_manager_v2/storage/core.cpp
     kv_cache_manager_v2/evictionController.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
     {