From 41340e120691c20f2899221d9a2b3944389af03b Mon Sep 17 00:00:00 2001 From: Chulian Zhang <851104+zhangcl@users.noreply.github.com> Date: Fri, 21 Aug 2026 23:03:23 -0700 Subject: [PATCH 1/4] [None][feat] Add locality domain runtime and bindings C++ locality domain utilities, the pluggable allocator and the runtime bindings. The Python layer that consumes them lands separately. Signed-off-by: Chulian Zhang <851104+zhangcl@users.noreply.github.com> --- .../nanobind/runtime/bindings.cpp | 104 ++ cpp/tensorrt_llm/runtime/CMakeLists.txt | 3 + .../localityDomainResourceConfig.h | 73 ++ .../locality_domain/locality_domain_utils.cpp | 1117 +++++++++++++++++ .../locality_domain/locality_domain_utils.h | 83 ++ cpp/tensorrt_llm/thop/CMakeLists.txt | 4 +- .../thop/localityDomainAllocator.cpp | 110 ++ cpp/tests/unit_tests/runtime/CMakeLists.txt | 2 + .../localityDomainPublicConfigTest.cpp | 93 ++ .../unit_tests/runtime/localizationTest.cu | 522 ++++++++ 10 files changed, 2110 insertions(+), 1 deletion(-) create mode 100644 cpp/tensorrt_llm/runtime/locality_domain/localityDomainResourceConfig.h create mode 100644 cpp/tensorrt_llm/runtime/locality_domain/locality_domain_utils.cpp create mode 100644 cpp/tensorrt_llm/runtime/locality_domain/locality_domain_utils.h create mode 100644 cpp/tensorrt_llm/thop/localityDomainAllocator.cpp create mode 100644 cpp/tests/unit_tests/runtime/localityDomainPublicConfigTest.cpp create mode 100644 cpp/tests/unit_tests/runtime/localizationTest.cu diff --git a/cpp/tensorrt_llm/nanobind/runtime/bindings.cpp b/cpp/tensorrt_llm/nanobind/runtime/bindings.cpp index eec3cd79bac1..85e38c555969 100644 --- a/cpp/tensorrt_llm/nanobind/runtime/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/runtime/bindings.cpp @@ -36,6 +36,7 @@ #include "tensorrt_llm/runtime/iGptDecoderBatched.h" #include "tensorrt_llm/runtime/iTensor.h" #include "tensorrt_llm/runtime/ipcUtils.h" +#include "tensorrt_llm/runtime/locality_domain/locality_domain_utils.h" #include "tensorrt_llm/runtime/lookaheadBuffers.h" #include "tensorrt_llm/runtime/loraCache.h" #include "tensorrt_llm/runtime/mcastGPUBuffer.h" @@ -53,6 +54,7 @@ #include #include #include +#include #include #include #include @@ -344,6 +346,108 @@ void initBindings(nb::module_& m) .value("ONESHOT", tensorrt_llm::kernels::AllReduceStrategyType::ONESHOT) .value("TWOSHOT", tensorrt_llm::kernels::AllReduceStrategyType::TWOSHOT); + // LOCALITY_DOMAIN Localization Handle bindings + nb::class_(m, "LocalizationHandle") + .def(nb::init<>(), nb::call_guard()) + .def("supports_localization", &tensorrt_llm::locality_domain::LocalizationHandle::supportsLocalization, + nb::call_guard()) + .def("supports_memory_localization", + &tensorrt_llm::locality_domain::LocalizationHandle::supportsMemoryLocalization, + nb::call_guard()) + .def("supports_compute_localization", + &tensorrt_llm::locality_domain::LocalizationHandle::supportsComputeLocalization, + nb::call_guard()) + .def( + "locality_domain_malloc", + [](tensorrt_llm::locality_domain::LocalizationHandle& self, size_t size, int localityDomainId) -> uintptr_t + { + void* ptr = nullptr; + self.localityDomainMalloc(&ptr, size, localityDomainId); + return reinterpret_cast(ptr); + }, + nb::arg("size"), nb::arg("locality_domain_id"), + "Allocate LOCALITY_DOMAIN localized memory and return pointer as integer address", + nb::call_guard()) + .def( + "locality_domain_free", + [](tensorrt_llm::locality_domain::LocalizationHandle& self, uintptr_t ptr) + { self.localityDomainFree(reinterpret_cast(ptr)); }, + nb::arg("ptr"), "Free LOCALITY_DOMAIN localized memory from integer address", + nb::call_guard()) + .def( + "create_localized_allocation_handle", + [](tensorrt_llm::locality_domain::LocalizationHandle& self, size_t size, int localityDomainId, + unsigned int requestedHandleTypes, bool gpuDirectRDMACapable, + std::optional usage) -> uintptr_t + { + CUmemGenericAllocationHandle const handle = usage.has_value() + ? self.createLocalizedAllocationHandle( + size, localityDomainId, requestedHandleTypes, gpuDirectRDMACapable, *usage) + : self.createLocalizedAllocationHandle( + size, localityDomainId, requestedHandleTypes, gpuDirectRDMACapable); + return static_cast(handle); + }, + nb::arg("size"), nb::arg("locality_domain_id"), nb::arg("requested_handle_types"), + nb::arg("gpu_direct_rdma_capable"), nb::arg("usage") = nb::none(), + "Create LOCALITY_DOMAIN localized generic allocation handle and return it as an integer", + nb::call_guard()) + .def( + "try_create_localized_allocation_handle", + [](tensorrt_llm::locality_domain::LocalizationHandle& self, size_t size, int localityDomainId, + unsigned int requestedHandleTypes, bool gpuDirectRDMACapable, + unsigned int usage) -> std::pair + { + CUmemGenericAllocationHandle handle{}; + CUresult const result = self.tryCreateLocalizedAllocationHandle( + &handle, size, localityDomainId, requestedHandleTypes, gpuDirectRDMACapable, usage); + return {static_cast(result), static_cast(handle)}; + }, + nb::arg("size"), nb::arg("locality_domain_id"), nb::arg("requested_handle_types"), + nb::arg("gpu_direct_rdma_capable"), nb::arg("usage"), + "Try to create a localized allocation and return (CUresult, handle)", + nb::call_guard()) + .def("get_localized_allocation_granularity", + &tensorrt_llm::locality_domain::LocalizationHandle::getLocalizedAllocationGranularity, + nb::arg("locality_domain_id"), nb::arg("requested_handle_types"), nb::arg("gpu_direct_rdma_capable"), + nb::arg("usage"), "Get minimum allocation granularity for a localized VMM allocation", + nb::call_guard()) + .def( + "try_get_localized_allocation_granularity", + [](tensorrt_llm::locality_domain::LocalizationHandle& self, int localityDomainId, + unsigned int requestedHandleTypes, bool gpuDirectRDMACapable, + unsigned int usage) -> std::pair + { + size_t granularity{}; + CUresult const result = self.tryGetLocalizedAllocationGranularity( + &granularity, localityDomainId, requestedHandleTypes, gpuDirectRDMACapable, usage); + return {static_cast(result), granularity}; + }, + nb::arg("locality_domain_id"), nb::arg("requested_handle_types"), nb::arg("gpu_direct_rdma_capable"), + nb::arg("usage"), "Try to get localized VMM granularity and return (CUresult, granularity)", + nb::call_guard()) + .def( + "create_localized_stream", + [](tensorrt_llm::locality_domain::LocalizationHandle& self, int localityDomainId) -> uintptr_t + { + CUstream stream = self.createLocalizedStream(localityDomainId); + return reinterpret_cast(stream); + }, + nb::arg("locality_domain_id"), + "Get a process-lifetime cached LOCALITY_DOMAIN localized stream as an integer address; callers must not " + "destroy it", + nb::call_guard()) + .def("get_locality_domain_compute_sm_counts", + &tensorrt_llm::locality_domain::LocalizationHandle::getLocalityDomainComputeSmCounts, + nb::arg("locality_domain_id"), + "Get (localized partition SM count, full-device SM count), or (0, 0) when unavailable", + nb::call_guard()) + .def( + "get_reserved_remainder_stream", + [](tensorrt_llm::locality_domain::LocalizationHandle& self) -> uintptr_t + { return reinterpret_cast(self.getReservedRemainderStream()); }, + "Get the borrowed process-lifetime remainder Green Context stream, or 0 when unavailable", + nb::call_guard()); + // Initialize MoeLoadBalancer bindings initMoeBindings(m); // Initialize HostFunc bindings diff --git a/cpp/tensorrt_llm/runtime/CMakeLists.txt b/cpp/tensorrt_llm/runtime/CMakeLists.txt index 11a9391c0e69..f7ca019a75dd 100644 --- a/cpp/tensorrt_llm/runtime/CMakeLists.txt +++ b/cpp/tensorrt_llm/runtime/CMakeLists.txt @@ -20,6 +20,7 @@ set(SRCS utils/runtimeUtils.cpp utils/debugUtils.cu utils/speculativeChoicesUtils.cpp + locality_domain/locality_domain_utils.cpp bufferManager.cpp cudaMemPool.cpp decodingLayerWorkspace.cpp @@ -76,6 +77,8 @@ set_property(TARGET runtime_src PROPERTY POSITION_INDEPENDENT_CODE ON) set_property(TARGET runtime_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON) add_cuda_architectures(runtime_src 89) +target_link_libraries(runtime_src PUBLIC ${CUDA_NVML_LIB}) + target_include_directories(runtime_src PRIVATE ${MPI_C_INCLUDE_DIRS}) if(ENABLE_MULTI_DEVICE) diff --git a/cpp/tensorrt_llm/runtime/locality_domain/localityDomainResourceConfig.h b/cpp/tensorrt_llm/runtime/locality_domain/localityDomainResourceConfig.h new file mode 100644 index 000000000000..07962f586414 --- /dev/null +++ b/cpp/tensorrt_llm/runtime/locality_domain/localityDomainResourceConfig.h @@ -0,0 +1,73 @@ +/* + * 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 + +#include + +#include + +namespace tensorrt_llm::locality_domain::detail +{ + +constexpr int kLocalityDomainCount = 2; + +#if CUDA_VERSION >= 13040 + +using SmResourceGroupParams = std::array; + +constexpr bool isBalancedSmCountValid(unsigned int totalSmCount) +{ + constexpr unsigned int kSmCountAlignment = 2; + constexpr unsigned int kMinimumTotalSmCount = kSmCountAlignment * static_cast(kLocalityDomainCount); + return totalSmCount >= kMinimumTotalSmCount && (totalSmCount % kMinimumTotalSmCount) == 0; +} + +constexpr bool isStrictSplitCountValid( + unsigned int totalSmCount, unsigned int localityDomainSmCount, unsigned int remainderSmCount) +{ + return localityDomainSmCount > 0 + && localityDomainSmCount <= totalSmCount / static_cast(kLocalityDomainCount) + && remainderSmCount == totalSmCount - localityDomainSmCount * static_cast(kLocalityDomainCount); +} + +inline SmResourceGroupParams makeStrictSmResourceGroupParams() +{ + SmResourceGroupParams groupParams{}; + for (int localityDomainId = 0; localityDomainId < kLocalityDomainCount; ++localityDomainId) + { + groupParams[localityDomainId].flags = CU_DEV_SM_RESOURCE_GROUP_LOCALITY_DOMAIN_ID; + groupParams[localityDomainId].localityDomainId = static_cast(localityDomainId); + } + return groupParams; +} + +inline SmResourceGroupParams makeBalancedSmResourceGroupParams(unsigned int totalSmCount) +{ + SmResourceGroupParams groupParams = makeStrictSmResourceGroupParams(); + unsigned int const smCountPerLocalityDomain = totalSmCount / static_cast(kLocalityDomainCount); + for (auto& params : groupParams) + { + params.smCount = smCountPerLocalityDomain; + params.flags |= CU_DEV_SM_RESOURCE_GROUP_BACKFILL; + } + return groupParams; +} + +#endif // CUDA_VERSION >= 13040 + +} // namespace tensorrt_llm::locality_domain::detail diff --git a/cpp/tensorrt_llm/runtime/locality_domain/locality_domain_utils.cpp b/cpp/tensorrt_llm/runtime/locality_domain/locality_domain_utils.cpp new file mode 100644 index 000000000000..398871b82bcc --- /dev/null +++ b/cpp/tensorrt_llm/runtime/locality_domain/locality_domain_utils.cpp @@ -0,0 +1,1117 @@ +/* + * 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 "tensorrt_llm/runtime/locality_domain/locality_domain_utils.h" + +#include "tensorrt_llm/common/cudaDriverWrapper.h" +#include "tensorrt_llm/common/logger.h" +#include "tensorrt_llm/runtime/locality_domain/localityDomainResourceConfig.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::locality_domain +{ + +namespace +{ + +#if CUDA_VERSION >= 13040 +constexpr int kLocalityDomainApiVersion = 13'040; +constexpr int kLocalityDomainCount = detail::kLocalityDomainCount; +#endif + +enum class LocalityDomainStreamCreateMethod +{ + kStrict, + kBalanced, +}; + +LocalityDomainStreamCreateMethod parseStreamCreateMethod() +{ + char const* const value = std::getenv("TLLM_LOCALITY_DOMAIN_STREAM_CREATE_METHOD"); + if (value == nullptr) + { + return LocalityDomainStreamCreateMethod::kStrict; + } + + std::string const method{value}; + if (method == "Balanced" || method == "balanced" || method == "BALANCED") + { + return LocalityDomainStreamCreateMethod::kBalanced; + } + + if (method == "GreenContext" || method == "greencontext" || method == "green" || method == "locality_domain" + || method == "LOCALITY_DOMAIN" || method == "3-part-gc" || method == "strict") + { + return LocalityDomainStreamCreateMethod::kStrict; + } + + TLLM_LOG_WARNING( + "[Localization] Unknown TLLM_LOCALITY_DOMAIN_STREAM_CREATE_METHOD=%s; using public strict Green Context split", + value); + return LocalityDomainStreamCreateMethod::kStrict; +} + +struct InstanceKey +{ + int device{}; + CUcontext context{}; +}; + +struct InstanceKeyLess +{ + bool operator()(InstanceKey const& lhs, InstanceKey const& rhs) const + { + if (lhs.device != rhs.device) + { + return lhs.device < rhs.device; + } + return std::less{}(lhs.context, rhs.context); + } +}; + +#if CUDA_VERSION >= 13040 + +struct AllocationKey +{ + int device{}; + CUcontext context{}; + CUdeviceptr pointer{}; +}; + +struct AllocationKeyLess +{ + bool operator()(AllocationKey const& lhs, AllocationKey const& rhs) const + { + if (lhs.device != rhs.device) + { + return lhs.device < rhs.device; + } + if (lhs.context != rhs.context) + { + return std::less{}(lhs.context, rhs.context); + } + return lhs.pointer < rhs.pointer; + } +}; + +struct VmmAllocation +{ + size_t alignedSize{}; + int device{}; + CUcontext context{}; + bool mapped{}; +}; + +std::map& getVmmAllocations() +{ + static auto* const allocations = new std::map; + return *allocations; +} + +std::mutex& getVmmAllocationMutex() +{ + static auto* const mutex = new std::mutex; + return *mutex; +} + +#endif + +CUresult queryCurrentDeviceAndContext(int* device, CUcontext* context) +{ + if (device == nullptr || context == nullptr) + { + return CUDA_ERROR_INVALID_VALUE; + } + + cudaError_t const runtimeResult = cudaFree(nullptr); + if (runtimeResult != cudaSuccess) + { + TLLM_LOG_WARNING( + "[Localization] Failed to initialize the current CUDA context: %s", cudaGetErrorString(runtimeResult)); + return CUDA_ERROR_INVALID_CONTEXT; + } + + int runtimeDevice{}; + cudaError_t const deviceResult = cudaGetDevice(&runtimeDevice); + if (deviceResult != cudaSuccess) + { + TLLM_LOG_WARNING( + "[Localization] Failed to query the current CUDA device: %s", cudaGetErrorString(deviceResult)); + return CUDA_ERROR_INVALID_DEVICE; + } + + CUcontext currentContext{}; + CUresult const contextResult = cuCtxGetCurrent(¤tContext); + if (contextResult != CUDA_SUCCESS) + { + return contextResult; + } + if (currentContext == nullptr) + { + return CUDA_ERROR_INVALID_CONTEXT; + } + + *device = runtimeDevice; + *context = currentContext; + return CUDA_SUCCESS; +} + +#if CUDA_VERSION >= 13040 + +template +bool loadDriverProc(char const* name, Proc* proc) +{ + void* address{}; + CUdriverProcAddressQueryResult queryResult{}; + CUresult const result + = cuGetProcAddress(name, &address, kLocalityDomainApiVersion, CU_GET_PROC_ADDRESS_DEFAULT, &queryResult); + if (result != CUDA_SUCCESS || queryResult != CU_GET_PROC_ADDRESS_SUCCESS || address == nullptr) + { + TLLM_LOG_WARNING("[Localization] CUDA 13.4 driver entry point %s is unavailable (result=%d, status=%d)", name, + static_cast(result), static_cast(queryResult)); + *proc = nullptr; + return false; + } + + *proc = reinterpret_cast(address); + return true; +} + +class GreenContextApi +{ +public: + using DeviceGetDevResource = CUresult(CUDAAPI*)(CUdevice device, CUdevResource* resource, CUdevResourceType type); + using DevSmResourceSplit + = CUresult(CUDAAPI*)(CUdevResource* result, unsigned int nbGroups, CUdevResource const* input, + CUdevResource* remainder, unsigned int flags, CU_DEV_SM_RESOURCE_GROUP_PARAMS* groupParams); + using DevResourceGenerateDesc + = CUresult(CUDAAPI*)(CUdevResourceDesc* desc, CUdevResource* resources, unsigned int nbResources); + using GreenCtxCreate + = CUresult(CUDAAPI*)(CUgreenCtx* greenContext, CUdevResourceDesc desc, CUdevice device, unsigned int flags); + using GreenCtxDestroy = CUresult(CUDAAPI*)(CUgreenCtx greenContext); + using GreenCtxStreamCreate + = CUresult(CUDAAPI*)(CUstream* stream, CUgreenCtx greenContext, unsigned int flags, int priority); + + bool load() + { + return loadDriverProc("cuDeviceGetDevResource", &deviceGetDevResource) + && loadDriverProc("cuDevSmResourceSplit", &devSmResourceSplit) + && loadDriverProc("cuDevResourceGenerateDesc", &devResourceGenerateDesc) + && loadDriverProc("cuGreenCtxCreate", &greenCtxCreate) + && loadDriverProc("cuGreenCtxDestroy", &greenCtxDestroy) + && loadDriverProc("cuGreenCtxStreamCreate", &greenCtxStreamCreate); + } + + DeviceGetDevResource deviceGetDevResource{}; + DevSmResourceSplit devSmResourceSplit{}; + DevResourceGenerateDesc devResourceGenerateDesc{}; + GreenCtxCreate greenCtxCreate{}; + GreenCtxDestroy greenCtxDestroy{}; + GreenCtxStreamCreate greenCtxStreamCreate{}; +}; + +class GreenContextPartitions +{ +public: + GreenContextPartitions() = default; + GreenContextPartitions(GreenContextPartitions const&) = delete; + GreenContextPartitions& operator=(GreenContextPartitions const&) = delete; + + ~GreenContextPartitions() + { + reset(); + } + + bool initialize(CUdevice device, LocalityDomainStreamCreateMethod method, unsigned int localityDomainSmCount) + { + mMethod = method; + if (!mApi.load()) + { + return false; + } + + CUdevResource fullResource{}; + CUresult result = mApi.deviceGetDevResource(device, &fullResource, CU_DEV_RESOURCE_TYPE_SM); + if (result != CUDA_SUCCESS) + { + TLLM_LOG_WARNING("[Localization] cuDeviceGetDevResource failed (result=%d)", static_cast(result)); + return false; + } + if (fullResource.type != CU_DEV_RESOURCE_TYPE_SM || fullResource.sm.smCount == 0) + { + TLLM_LOG_WARNING("[Localization] Device returned an invalid SM resource"); + return false; + } + + detail::SmResourceGroupParams groupParams = detail::makeStrictSmResourceGroupParams(); + if (method == LocalityDomainStreamCreateMethod::kBalanced) + { + if (!detail::isBalancedSmCountValid(fullResource.sm.smCount)) + { + TLLM_LOG_WARNING( + "[Localization] Balanced public split requires an even per-group SM count, got " + "total SM count %u", + fullResource.sm.smCount); + return false; + } + groupParams = detail::makeBalancedSmResourceGroupParams(fullResource.sm.smCount); + } + + result = mApi.devSmResourceSplit(mLocalizedResources.data(), kLocalityDomainCount, &fullResource, + &mRemainderResource, + /*flags=*/0, groupParams.data()); + if (result != CUDA_SUCCESS) + { + TLLM_LOG_WARNING( + "[Localization] Public locality-domain SM split failed (result=%d)", static_cast(result)); + return false; + } + + for (int localityDomainId = 0; localityDomainId < kLocalityDomainCount; ++localityDomainId) + { + CUdevResource const& resource = mLocalizedResources[localityDomainId]; + if (resource.type != CU_DEV_RESOURCE_TYPE_SM || resource.sm.smCount == 0 + || resource.sm.localityDomainId != static_cast(localityDomainId) + || (resource.sm.flags & CU_DEV_SM_RESOURCE_GROUP_LOCALITY_DOMAIN_ID) == 0) + { + TLLM_LOG_WARNING( + "[Localization] Public split returned an invalid resource for locality domain%d", localityDomainId); + reset(); + return false; + } + + if (method == LocalityDomainStreamCreateMethod::kStrict && resource.sm.smCount != localityDomainSmCount) + { + TLLM_LOG_WARNING( + "[Localization] Strict public split returned %u SMs for locality domain%d, expected " + "locality-domain " + "attribute value %u", + resource.sm.smCount, localityDomainId, localityDomainSmCount); + reset(); + return false; + } + + if (method == LocalityDomainStreamCreateMethod::kBalanced + && resource.sm.smCount != fullResource.sm.smCount / static_cast(kLocalityDomainCount)) + { + TLLM_LOG_WARNING( + "[Localization] Balanced public split returned %u SMs for locality domain%d, expected %u", + resource.sm.smCount, localityDomainId, + fullResource.sm.smCount / static_cast(kLocalityDomainCount)); + reset(); + return false; + } + } + + if (method == LocalityDomainStreamCreateMethod::kStrict) + { + if (!detail::isStrictSplitCountValid(fullResource.sm.smCount, localityDomainSmCount, getRemainderSmCount())) + { + TLLM_LOG_WARNING( + "[Localization] Strict public split counts do not match two complete locality domains " + "(total=%u, per-domain=%u, remainder=%u)", + fullResource.sm.smCount, localityDomainSmCount, getRemainderSmCount()); + reset(); + return false; + } + } + else if (getRemainderSmCount() != 0) + { + TLLM_LOG_WARNING( + "[Localization] Balanced public split unexpectedly left %u remainder SMs", getRemainderSmCount()); + reset(); + return false; + } + + TLLM_LOG_INFO( + "[Localization] Public %s split: total=%u SM, locality domain 0=%u SM, locality domain 1=%u SM, " + "remainder=%u SM", + method == LocalityDomainStreamCreateMethod::kStrict ? "strict" : "balanced", fullResource.sm.smCount, + mLocalizedResources[0].sm.smCount, mLocalizedResources[1].sm.smCount, getRemainderSmCount()); + + for (int localityDomainId = 0; localityDomainId < kLocalityDomainCount; ++localityDomainId) + { + result = mApi.devResourceGenerateDesc( + &mLocalizedDescriptors[localityDomainId], &mLocalizedResources[localityDomainId], /*nbResources=*/1); + if (result == CUDA_SUCCESS) + { + result = mApi.greenCtxCreate(&mGreenContexts[localityDomainId], mLocalizedDescriptors[localityDomainId], + device, CU_GREEN_CTX_DEFAULT_STREAM); + } + if (result != CUDA_SUCCESS) + { + TLLM_LOG_WARNING("[Localization] Failed to create Green Context for locality domain%d (result=%d)", + localityDomainId, static_cast(result)); + reset(); + return false; + } + } + + if (method == LocalityDomainStreamCreateMethod::kStrict && getRemainderSmCount() > 0) + { + result = mApi.devResourceGenerateDesc(&mRemainderDescriptor, &mRemainderResource, /*nbResources=*/1); + if (result == CUDA_SUCCESS) + { + result = mApi.greenCtxCreate( + &mRemainderGreenContext, mRemainderDescriptor, device, CU_GREEN_CTX_DEFAULT_STREAM); + } + if (result == CUDA_SUCCESS) + { + result = mApi.greenCtxStreamCreate( + &mRemainderStream, mRemainderGreenContext, CU_STREAM_NON_BLOCKING, /*priority=*/0); + } + if (result != CUDA_SUCCESS) + { + TLLM_LOG_WARNING( + "[Localization] Failed to create strict split remainder Green Context/stream (result=%d)", + static_cast(result)); + reset(); + return false; + } + } + + mSupported = true; + return true; + } + + CUresult createStream(CUstream* stream, int localityDomainId) + { + if (stream == nullptr || localityDomainId < 0 || localityDomainId >= kLocalityDomainCount) + { + return CUDA_ERROR_INVALID_VALUE; + } + if (!mSupported) + { + return CUDA_ERROR_NOT_SUPPORTED; + } + + std::lock_guard const lock{mStreamMutex}; + if (mLocalizedStreams[localityDomainId] == nullptr) + { + CUstream newStream{}; + CUresult const result = mApi.greenCtxStreamCreate( + &newStream, mGreenContexts[localityDomainId], CU_STREAM_NON_BLOCKING, /*priority=*/0); + if (result != CUDA_SUCCESS) + { + if (newStream != nullptr) + { + static_cast(cuStreamDestroy(newStream)); + } + return result; + } + mLocalizedStreams[localityDomainId] = newStream; + } + *stream = mLocalizedStreams[localityDomainId]; + return CUDA_SUCCESS; + } + + CUstream getRemainderStream() const + { + if (!mSupported || mMethod != LocalityDomainStreamCreateMethod::kStrict) + { + return nullptr; + } + return mRemainderStream; + } + + std::pair getSmCounts(int localityDomainId) const noexcept + { + if (!mSupported || localityDomainId < 0 || localityDomainId >= kLocalityDomainCount) + { + return {}; + } + + CUdevResource const& localizedResource = mLocalizedResources[localityDomainId]; + if (localizedResource.type != CU_DEV_RESOURCE_TYPE_SM || localizedResource.sm.smCount == 0) + { + return {}; + } + + unsigned int totalSmCount = getRemainderSmCount(); + for (auto const& resource : mLocalizedResources) + { + if (resource.type != CU_DEV_RESOURCE_TYPE_SM || resource.sm.smCount == 0) + { + return {}; + } + totalSmCount += resource.sm.smCount; + } + return {localizedResource.sm.smCount, totalSmCount}; + } + +private: + unsigned int getRemainderSmCount() const + { + return mRemainderResource.type == CU_DEV_RESOURCE_TYPE_SM ? mRemainderResource.sm.smCount : 0; + } + + void reset() noexcept + { + for (auto& stream : mLocalizedStreams) + { + if (stream != nullptr) + { + CUresult const result = cuStreamDestroy(stream); + if (result != CUDA_SUCCESS) + { + TLLM_LOG_WARNING( + "[Localization] Failed to destroy localized stream (result=%d)", static_cast(result)); + } + stream = nullptr; + } + } + + if (mRemainderStream != nullptr && mApi.greenCtxDestroy != nullptr) + { + CUresult const result = cuStreamDestroy(mRemainderStream); + if (result != CUDA_SUCCESS) + { + TLLM_LOG_WARNING( + "[Localization] Failed to destroy remainder stream (result=%d)", static_cast(result)); + } + mRemainderStream = nullptr; + } + + if (mRemainderGreenContext != nullptr && mApi.greenCtxDestroy != nullptr) + { + CUresult const result = mApi.greenCtxDestroy(mRemainderGreenContext); + if (result != CUDA_SUCCESS) + { + TLLM_LOG_WARNING( + "[Localization] Failed to destroy remainder Green Context (result=%d)", static_cast(result)); + } + mRemainderGreenContext = nullptr; + } + + if (mApi.greenCtxDestroy != nullptr) + { + for (auto& greenContext : mGreenContexts) + { + if (greenContext != nullptr) + { + CUresult const result = mApi.greenCtxDestroy(greenContext); + if (result != CUDA_SUCCESS) + { + TLLM_LOG_WARNING( + "[Localization] Failed to destroy Green Context (result=%d)", static_cast(result)); + } + greenContext = nullptr; + } + } + } + mSupported = false; + } + + GreenContextApi mApi; + LocalityDomainStreamCreateMethod mMethod{LocalityDomainStreamCreateMethod::kStrict}; + std::array mLocalizedResources{}; + std::array mLocalizedDescriptors{}; + std::array mGreenContexts{}; + std::array mLocalizedStreams{}; + CUdevResource mRemainderResource{}; + CUdevResourceDesc mRemainderDescriptor{}; + CUgreenCtx mRemainderGreenContext{}; + CUstream mRemainderStream{}; + bool mSupported{}; + std::mutex mStreamMutex; +}; + +#endif // CUDA_VERSION >= 13040 + +} // namespace + +class Localization +{ +public: + Localization(int device, CUcontext context) + : mDevice{device} + , mContext{context} + , mStreamCreateMethod{parseStreamCreateMethod()} + { + initialize(); + } + + bool supportsMemoryLocalization() const noexcept + { + return mMemoryLocalizationSupported; + } + + bool supportsComputeLocalization() const noexcept + { + return mComputeLocalizationSupported; + } + + bool supportsLocalization() const noexcept + { + return supportsMemoryLocalization() && supportsComputeLocalization(); + } + + CUresult localizedDeviceAlloc(void** localizedDevPtr, size_t size, int localityDomainId) noexcept + { + if (localizedDevPtr == nullptr || size == 0) + { + return CUDA_ERROR_INVALID_VALUE; + } + *localizedDevPtr = nullptr; + + if (localityDomainId == -1) + { + return cuMemAlloc(reinterpret_cast(localizedDevPtr), size); + } + +#if CUDA_VERSION >= 13040 + if (!mMemoryLocalizationSupported) + { + return CUDA_ERROR_NOT_SUPPORTED; + } + + CUmemAllocationProp prop{}; + CUresult result = makeAllocationProp(&prop, localityDomainId, CU_MEM_HANDLE_TYPE_NONE, + /*gpuDirectRDMACapable=*/false, /*usage=*/0); + if (result != CUDA_SUCCESS) + { + return result; + } + + size_t granularity{}; + result = cuMemGetAllocationGranularity(&granularity, &prop, CU_MEM_ALLOC_GRANULARITY_MINIMUM); + if (result != CUDA_SUCCESS) + { + return result; + } + if (granularity == 0 || size > std::numeric_limits::max() - (granularity - 1)) + { + return CUDA_ERROR_INVALID_VALUE; + } + size_t const alignedSize = ((size + granularity - 1) / granularity) * granularity; + + CUdeviceptr address{}; + CUmemGenericAllocationHandle allocationHandle{}; + bool mapped = false; + + result = cuMemAddressReserve(&address, alignedSize, /*alignment=*/0, /*addr=*/0, /*flags=*/0); + if (result == CUDA_SUCCESS) + { + result = cuMemCreate(&allocationHandle, alignedSize, &prop, /*flags=*/0); + } + if (result == CUDA_SUCCESS) + { + result = cuMemMap(address, alignedSize, /*offset=*/0, allocationHandle, /*flags=*/0); + mapped = result == CUDA_SUCCESS; + } + if (result == CUDA_SUCCESS) + { + CUmemAccessDesc access{}; + access.location.type = CU_MEM_LOCATION_TYPE_DEVICE; + access.location.id = mDevice; + access.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; + result = cuMemSetAccess(address, alignedSize, &access, /*count=*/1); + } + if (result == CUDA_SUCCESS) + { + result = cuMemRelease(allocationHandle); + if (result == CUDA_SUCCESS) + { + allocationHandle = 0; + } + } + + if (result == CUDA_SUCCESS) + { + try + { + AllocationKey const key{mDevice, mContext, address}; + VmmAllocation const allocation{alignedSize, mDevice, mContext, mapped}; + std::lock_guard const lock{getVmmAllocationMutex()}; + bool const inserted = getVmmAllocations().emplace(key, allocation).second; + if (!inserted) + { + result = CUDA_ERROR_INVALID_VALUE; + } + } + catch (...) + { + result = CUDA_ERROR_OUT_OF_MEMORY; + } + } + + if (result != CUDA_SUCCESS) + { + if (mapped) + { + static_cast(cuMemUnmap(address, alignedSize)); + } + if (allocationHandle != 0) + { + static_cast(cuMemRelease(allocationHandle)); + } + if (address != 0) + { + static_cast(cuMemAddressFree(address, alignedSize)); + } + return result; + } + + *localizedDevPtr = reinterpret_cast(address); + return CUDA_SUCCESS; +#else + static_cast(localityDomainId); + return CUDA_ERROR_NOT_SUPPORTED; +#endif + } + + CUresult localizedDeviceFree(void* localizedDevPtr) + { + if (localizedDevPtr == nullptr) + { + return CUDA_SUCCESS; + } + +#if CUDA_VERSION >= 13040 + CUresult result = checkCurrentContext(); + if (result != CUDA_SUCCESS) + { + return result; + } + + CUdeviceptr const address = reinterpret_cast(localizedDevPtr); + AllocationKey const key{mDevice, mContext, address}; + std::lock_guard const lock{getVmmAllocationMutex()}; + auto& allocations = getVmmAllocations(); + auto allocationIt = allocations.find(key); + if (allocationIt == allocations.end()) + { + return cuMemFree(address); + } + + VmmAllocation& allocation = allocationIt->second; + if (allocation.device != mDevice || allocation.context != mContext) + { + return CUDA_ERROR_INVALID_CONTEXT; + } + if (allocation.mapped) + { + result = cuMemUnmap(address, allocation.alignedSize); + if (result != CUDA_SUCCESS) + { + return result; + } + allocation.mapped = false; + } + + result = cuMemAddressFree(address, allocation.alignedSize); + if (result == CUDA_SUCCESS) + { + allocations.erase(allocationIt); + } + return result; +#else + return cuMemFree(reinterpret_cast(localizedDevPtr)); +#endif + } + + CUresult tryCreateLocalizedAllocationHandle(CUmemGenericAllocationHandle* handle, size_t size, int localityDomainId, + unsigned int requestedHandleTypes, bool gpuDirectRDMACapable, unsigned int usage) noexcept + { + if (handle == nullptr || size == 0) + { + return CUDA_ERROR_INVALID_VALUE; + } + *handle = 0; + +#if CUDA_VERSION >= 13040 + CUmemAllocationProp prop{}; + CUresult result + = makeAllocationProp(&prop, localityDomainId, requestedHandleTypes, gpuDirectRDMACapable, usage); + if (result != CUDA_SUCCESS) + { + return result; + } + + size_t granularity{}; + result = cuMemGetAllocationGranularity(&granularity, &prop, CU_MEM_ALLOC_GRANULARITY_MINIMUM); + if (result != CUDA_SUCCESS) + { + return result; + } + if (granularity == 0 || (size % granularity) != 0) + { + return CUDA_ERROR_INVALID_VALUE; + } + return cuMemCreate(handle, size, &prop, /*flags=*/0); +#else + static_cast(localityDomainId); + static_cast(requestedHandleTypes); + static_cast(gpuDirectRDMACapable); + static_cast(usage); + return CUDA_ERROR_NOT_SUPPORTED; +#endif + } + + CUresult tryGetLocalizedAllocationGranularity(size_t* granularity, int localityDomainId, + unsigned int requestedHandleTypes, bool gpuDirectRDMACapable, unsigned int usage) noexcept + { + if (granularity == nullptr) + { + return CUDA_ERROR_INVALID_VALUE; + } + *granularity = 0; + +#if CUDA_VERSION >= 13040 + CUmemAllocationProp prop{}; + CUresult const result + = makeAllocationProp(&prop, localityDomainId, requestedHandleTypes, gpuDirectRDMACapable, usage); + if (result != CUDA_SUCCESS) + { + return result; + } + return cuMemGetAllocationGranularity(granularity, &prop, CU_MEM_ALLOC_GRANULARITY_MINIMUM); +#else + static_cast(localityDomainId); + static_cast(requestedHandleTypes); + static_cast(gpuDirectRDMACapable); + static_cast(usage); + return CUDA_ERROR_NOT_SUPPORTED; +#endif + } + + unsigned int getAutomaticAllocationUsage(bool gpuDirectRDMACapable) const noexcept + { +#if CUDA_VERSION >= 13040 + if (gpuDirectRDMACapable && !mDefaultRdmaSupportsLocalizedMemory) + { + return CU_MEM_CREATE_USAGE_GPU_DIRECT_RDMA_OVER_PCIE; + } +#else + static_cast(gpuDirectRDMACapable); +#endif + return 0; + } + + CUresult createLocalizedStream(CUstream* stream, int localityDomainId) + { + if (stream == nullptr) + { + return CUDA_ERROR_INVALID_VALUE; + } + *stream = nullptr; + +#if CUDA_VERSION >= 13040 + CUresult const contextResult = checkCurrentContext(); + if (contextResult != CUDA_SUCCESS) + { + return contextResult; + } + return mPartitions.createStream(stream, localityDomainId); +#else + static_cast(localityDomainId); + return CUDA_ERROR_NOT_SUPPORTED; +#endif + } + + CUstream getReservedRemainderStream() const noexcept + { +#if CUDA_VERSION >= 13040 + return mPartitions.getRemainderStream(); +#else + return nullptr; +#endif + } + + std::pair getComputeSmCounts(int localityDomainId) const noexcept + { +#if CUDA_VERSION >= 13040 + return mPartitions.getSmCounts(localityDomainId); +#else + static_cast(localityDomainId); + return {}; +#endif + } + + static Localization* getLocalization() + { + int device{}; + CUcontext context{}; + TLLM_CU_CHECK(queryCurrentDeviceAndContext(&device, &context)); + + static auto* const mutex = new std::mutex; + static auto* const localizations = new std::map; + + InstanceKey const key{device, context}; + std::lock_guard const lock{*mutex}; + auto const it = localizations->find(key); + if (it != localizations->end()) + { + return it->second; + } + + auto* const localization = new Localization(device, context); + localizations->emplace(key, localization); + return localization; + } + +private: + void initialize() + { +#if CUDA_VERSION >= 13040 + int localityDomainCount{}; + CUresult result + = cuDeviceGetAttribute(&localityDomainCount, CU_DEVICE_ATTRIBUTE_LOCALITY_DOMAIN_COUNT, mDevice); + if (result != CUDA_SUCCESS || localityDomainCount < kLocalityDomainCount) + { + TLLM_LOG_INFO( + "[Localization] Device %d does not expose two public locality domains (result=%d, " + "count=%d)", + mDevice, static_cast(result), localityDomainCount); + return; + } + mLocalityDomainCount = localityDomainCount; + mMemoryLocalizationSupported = true; + + int defaultRdmaSupported{}; + result = cuDeviceGetAttribute( + &defaultRdmaSupported, CU_DEVICE_ATTRIBUTE_GPU_DIRECT_RDMA_WITH_LOCALIZED_MEMORY_SUPPORTED, mDevice); + if (result == CUDA_SUCCESS) + { + mDefaultRdmaSupportsLocalizedMemory = defaultRdmaSupported != 0; + } + + int localityDomainSmCount{}; + result = cuDeviceGetAttribute( + &localityDomainSmCount, CU_DEVICE_ATTRIBUTE_LOCALITY_DOMAIN_MULTIPROCESSOR_COUNT, mDevice); + if (result != CUDA_SUCCESS || localityDomainSmCount <= 0) + { + TLLM_LOG_WARNING("[Localization] Failed to query locality-domain SM count (result=%d, count=%d)", + static_cast(result), localityDomainSmCount); + return; + } + + mComputeLocalizationSupported + = mPartitions.initialize(mDevice, mStreamCreateMethod, static_cast(localityDomainSmCount)); + if (!mComputeLocalizationSupported) + { + TLLM_LOG_WARNING( + "[Localization] Public localized VMM is available, but public Green Context " + "partitioning is unavailable"); + } +#else + TLLM_LOG_INFO("[Localization] Built with CUDA %d; public locality-domain support requires CUDA 13.4 headers", + CUDA_VERSION); +#endif + } + + CUresult checkCurrentContext() const noexcept + { + CUcontext currentContext{}; + CUresult result = cuCtxGetCurrent(¤tContext); + if (result != CUDA_SUCCESS) + { + return result; + } + if (currentContext != mContext) + { + return CUDA_ERROR_INVALID_CONTEXT; + } + + CUdevice currentDevice{}; + result = cuCtxGetDevice(¤tDevice); + if (result != CUDA_SUCCESS) + { + return result; + } + return currentDevice == mDevice ? CUDA_SUCCESS : CUDA_ERROR_INVALID_CONTEXT; + } + +#if CUDA_VERSION >= 13040 + CUresult makeAllocationProp(CUmemAllocationProp* prop, int localityDomainId, unsigned int requestedHandleTypes, + bool gpuDirectRDMACapable, unsigned int usage) const noexcept + { + if (prop == nullptr || localityDomainId < 0 || localityDomainId >= kLocalityDomainCount || mDevice < 0 + || mDevice > std::numeric_limits::max() + || usage > std::numeric_limits::max()) + { + return CUDA_ERROR_INVALID_VALUE; + } + if (!mMemoryLocalizationSupported) + { + return CUDA_ERROR_NOT_SUPPORTED; + } + if (localityDomainId >= mLocalityDomainCount) + { + return CUDA_ERROR_INVALID_VALUE; + } + + CUresult const contextResult = checkCurrentContext(); + if (contextResult != CUDA_SUCCESS) + { + return contextResult; + } + + *prop = {}; + prop->type = CU_MEM_ALLOCATION_TYPE_PINNED; + prop->requestedHandleTypes = static_cast(requestedHandleTypes); + prop->location.type = CU_MEM_LOCATION_TYPE_DEVICE_LOCALITY_DOMAIN; + prop->location.localized.deviceId = static_cast(mDevice); + prop->location.localized.localityDomainId = static_cast(localityDomainId); + prop->allocFlags.gpuDirectRDMACapable = gpuDirectRDMACapable ? 1 : 0; + prop->allocFlags.usage = static_cast(usage); + return CUDA_SUCCESS; + } +#endif + + int mDevice{}; + CUcontext mContext{}; + LocalityDomainStreamCreateMethod mStreamCreateMethod{LocalityDomainStreamCreateMethod::kStrict}; + bool mMemoryLocalizationSupported{}; + bool mComputeLocalizationSupported{}; +#if CUDA_VERSION >= 13040 + bool mDefaultRdmaSupportsLocalizedMemory{}; + int mLocalityDomainCount{}; + GreenContextPartitions mPartitions; +#endif +}; + +LocalizationHandle::LocalizationHandle() + : mImpl{Localization::getLocalization()} +{ +} + +LocalizationHandle::~LocalizationHandle() = default; + +LocalizationHandle::LocalizationHandle(LocalizationHandle&& other) noexcept + : mImpl{other.mImpl} +{ + other.mImpl = nullptr; +} + +LocalizationHandle& LocalizationHandle::operator=(LocalizationHandle&& other) noexcept +{ + if (this != &other) + { + mImpl = other.mImpl; + other.mImpl = nullptr; + } + return *this; +} + +bool LocalizationHandle::supportsLocalization() const +{ + return mImpl != nullptr && mImpl->supportsLocalization(); +} + +bool LocalizationHandle::supportsMemoryLocalization() const +{ + return mImpl != nullptr && mImpl->supportsMemoryLocalization(); +} + +bool LocalizationHandle::supportsComputeLocalization() const +{ + return mImpl != nullptr && mImpl->supportsComputeLocalization(); +} + +void LocalizationHandle::localityDomainMalloc(void** localizedDevPtr, size_t size, int localityDomainId) +{ + TLLM_CHECK_WITH_INFO(mImpl != nullptr, "Cannot use a moved-from LocalizationHandle"); + TLLM_CU_CHECK(mImpl->localizedDeviceAlloc(localizedDevPtr, size, localityDomainId)); +} + +void LocalizationHandle::localityDomainFree(void* localizedDevPtr) +{ + TLLM_CHECK_WITH_INFO(mImpl != nullptr, "Cannot use a moved-from LocalizationHandle"); + TLLM_CU_CHECK(mImpl->localizedDeviceFree(localizedDevPtr)); +} + +CUmemGenericAllocationHandle LocalizationHandle::createLocalizedAllocationHandle( + size_t size, int localityDomainId, unsigned int requestedHandleTypes, bool gpuDirectRDMACapable) +{ + TLLM_CHECK_WITH_INFO(mImpl != nullptr, "Cannot use a moved-from LocalizationHandle"); + return createLocalizedAllocationHandle(size, localityDomainId, requestedHandleTypes, gpuDirectRDMACapable, + mImpl->getAutomaticAllocationUsage(gpuDirectRDMACapable)); +} + +CUmemGenericAllocationHandle LocalizationHandle::createLocalizedAllocationHandle( + size_t size, int localityDomainId, unsigned int requestedHandleTypes, bool gpuDirectRDMACapable, unsigned int usage) +{ + CUmemGenericAllocationHandle handle{}; + TLLM_CU_CHECK(tryCreateLocalizedAllocationHandle( + &handle, size, localityDomainId, requestedHandleTypes, gpuDirectRDMACapable, usage)); + return handle; +} + +CUresult LocalizationHandle::tryCreateLocalizedAllocationHandle(CUmemGenericAllocationHandle* handle, size_t size, + int localityDomainId, unsigned int requestedHandleTypes, bool gpuDirectRDMACapable, unsigned int usage) noexcept +{ + if (mImpl == nullptr) + { + return CUDA_ERROR_INVALID_CONTEXT; + } + return mImpl->tryCreateLocalizedAllocationHandle( + handle, size, localityDomainId, requestedHandleTypes, gpuDirectRDMACapable, usage); +} + +size_t LocalizationHandle::getLocalizedAllocationGranularity( + int localityDomainId, unsigned int requestedHandleTypes, bool gpuDirectRDMACapable, unsigned int usage) +{ + size_t granularity{}; + TLLM_CU_CHECK(tryGetLocalizedAllocationGranularity( + &granularity, localityDomainId, requestedHandleTypes, gpuDirectRDMACapable, usage)); + return granularity; +} + +CUresult LocalizationHandle::tryGetLocalizedAllocationGranularity(size_t* granularity, int localityDomainId, + unsigned int requestedHandleTypes, bool gpuDirectRDMACapable, unsigned int usage) noexcept +{ + if (mImpl == nullptr) + { + return CUDA_ERROR_INVALID_CONTEXT; + } + return mImpl->tryGetLocalizedAllocationGranularity( + granularity, localityDomainId, requestedHandleTypes, gpuDirectRDMACapable, usage); +} + +CUstream LocalizationHandle::createLocalizedStream(int localityDomainId) +{ + TLLM_CHECK_WITH_INFO(mImpl != nullptr, "Cannot use a moved-from LocalizationHandle"); + CUstream stream{}; + TLLM_CU_CHECK(mImpl->createLocalizedStream(&stream, localityDomainId)); + return stream; +} + +std::pair LocalizationHandle::getLocalityDomainComputeSmCounts( + int localityDomainId) const noexcept +{ + return mImpl != nullptr ? mImpl->getComputeSmCounts(localityDomainId) : std::pair{}; +} + +CUstream LocalizationHandle::getReservedRemainderStream() +{ + TLLM_CHECK_WITH_INFO(mImpl != nullptr, "Cannot use a moved-from LocalizationHandle"); + return mImpl->getReservedRemainderStream(); +} + +} // namespace tensorrt_llm::locality_domain diff --git a/cpp/tensorrt_llm/runtime/locality_domain/locality_domain_utils.h b/cpp/tensorrt_llm/runtime/locality_domain/locality_domain_utils.h new file mode 100644 index 000000000000..6ddf2e815fad --- /dev/null +++ b/cpp/tensorrt_llm/runtime/locality_domain/locality_domain_utils.h @@ -0,0 +1,83 @@ +/* + * 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 + +#include + +#include +#include + +namespace tensorrt_llm +{ + +namespace locality_domain +{ + +class Localization; + +class LocalizationHandle +{ +public: + LocalizationHandle(); + ~LocalizationHandle(); + + // Delete copy constructor and copy assignment + LocalizationHandle(LocalizationHandle const&) = delete; + LocalizationHandle& operator=(LocalizationHandle const&) = delete; + + // Allow move constructor and move assignment + LocalizationHandle(LocalizationHandle&&) noexcept; + LocalizationHandle& operator=(LocalizationHandle&&) noexcept; + + //! Return whether both public compute and memory locality-domain APIs are usable. + bool supportsLocalization() const; + //! Return whether public locality-domain VMM allocation is usable. + bool supportsMemoryLocalization() const; + //! Return whether public locality-domain Green Context partitioning is usable. + bool supportsComputeLocalization() const; + + void localityDomainMalloc(void** localizedDevPtr, size_t size, int localityDomainId); + void localityDomainFree(void* localizedDevPtr); + + CUmemGenericAllocationHandle createLocalizedAllocationHandle( + size_t size, int localityDomainId, unsigned int requestedHandleTypes, bool gpuDirectRDMACapable); + CUmemGenericAllocationHandle createLocalizedAllocationHandle(size_t size, int localityDomainId, + unsigned int requestedHandleTypes, bool gpuDirectRDMACapable, unsigned int usage); + CUresult tryCreateLocalizedAllocationHandle(CUmemGenericAllocationHandle* handle, size_t size, int localityDomainId, + unsigned int requestedHandleTypes, bool gpuDirectRDMACapable, unsigned int usage) noexcept; + + size_t getLocalizedAllocationGranularity( + int localityDomainId, unsigned int requestedHandleTypes, bool gpuDirectRDMACapable, unsigned int usage); + CUresult tryGetLocalizedAllocationGranularity(size_t* granularity, int localityDomainId, + unsigned int requestedHandleTypes, bool gpuDirectRDMACapable, unsigned int usage) noexcept; + + //! Return a process-lifetime cached stream owned by the localization singleton. + //! Callers must not destroy the returned stream. + CUstream createLocalizedStream(int localityDomainId); + //! Return (localized partition SM count, full-device SM count), or (0, 0) when unavailable. + std::pair getLocalityDomainComputeSmCounts(int localityDomainId) const noexcept; + //! Return a borrowed process-lifetime remainder stream, or nullptr. The caller must not destroy it. + CUstream getReservedRemainderStream(); + +private: + Localization* mImpl; +}; + +} // namespace locality_domain + +} // namespace tensorrt_llm diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index 7af68a90f1c9..d31b3adbc1d6 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -126,6 +126,7 @@ add_library( relativeAttentionBiasOp.cpp dsv3RouterGemmOp.cpp customMoeRoutingOp.cpp + localityDomainAllocator.cpp mamba2MTPSSMCacheOp.cpp selectiveScanOp.cpp userbuffersFinalizeOp.cpp @@ -174,7 +175,8 @@ endif() if(ENABLE_MULTI_DEVICE) target_include_directories(th_common PUBLIC ${MPI_C_INCLUDE_DIRS}) - target_link_libraries(th_common PRIVATE ${MPI_C_LIBRARIES} ${NCCL_LIB}) + target_link_libraries(th_common PRIVATE ${MPI_C_LIBRARIES} ${NCCL_LIB} + CUDA::nvml) endif() if(NOT WIN32) diff --git a/cpp/tensorrt_llm/thop/localityDomainAllocator.cpp b/cpp/tensorrt_llm/thop/localityDomainAllocator.cpp new file mode 100644 index 000000000000..22b259a5f618 --- /dev/null +++ b/cpp/tensorrt_llm/thop/localityDomainAllocator.cpp @@ -0,0 +1,110 @@ +/* + * 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 "tensorrt_llm/common/logger.h" +#include "tensorrt_llm/runtime/locality_domain/locality_domain_utils.h" + +#include + +#include +#include + +#if defined(_WIN32) +#define TLLM_LOCALITY_DOMAIN_ALLOCATOR_EXPORT __declspec(dllexport) +#else +#define TLLM_LOCALITY_DOMAIN_ALLOCATOR_EXPORT __attribute__((visibility("default"))) +#endif + +namespace torch_ext +{ + +void* localityDomainLocalizationAlloc(size_t size, int device, void* stream, int localityDomainId) noexcept +{ + try + { + TLLM_LOG_DEBUG("localityDomainLocalizationAlloc: allocating %zu bytes memory for localityDomainId=%d", size, + localityDomainId); + std::optional deviceGuard; + if (device >= 0) + { + deviceGuard.emplace(static_cast(device)); + } + auto handle = tensorrt_llm::locality_domain::LocalizationHandle(); + void* outputPtr = nullptr; + handle.localityDomainMalloc(&outputPtr, size, localityDomainId); + return outputPtr; + } + catch (std::exception const& exception) + { + TLLM_LOG_EXCEPTION(exception); + } + catch (...) + { + TLLM_LOG_ERROR("Unknown exception thrown allocating locality domain-localized memory"); + } + return nullptr; +} + +void localityDomainLocalizationFree(void* ptr, size_t size, int device, void* stream, int localityDomainId) noexcept +{ + try + { + TLLM_LOG_DEBUG( + "localityDomainLocalizationFree: free %zu bytes memory for localityDomainId=%d", size, localityDomainId); + std::optional deviceGuard; + if (device >= 0) + { + deviceGuard.emplace(static_cast(device)); + } + auto handle = tensorrt_llm::locality_domain::LocalizationHandle(); + handle.localityDomainFree(ptr); + } + catch (std::exception const& exception) + { + TLLM_LOG_EXCEPTION(exception); + } + catch (...) + { + TLLM_LOG_ERROR("Unknown exception thrown freeing locality domain-localized memory"); + } +} + +} // namespace torch_ext + +extern "C" TLLM_LOCALITY_DOMAIN_ALLOCATOR_EXPORT void* trtllm_locality_domain0_alloc( + size_t size, int device, void* stream) noexcept +{ + return torch_ext::localityDomainLocalizationAlloc(size, device, stream, 0); +} + +extern "C" TLLM_LOCALITY_DOMAIN_ALLOCATOR_EXPORT void* trtllm_locality_domain1_alloc( + size_t size, int device, void* stream) noexcept +{ + return torch_ext::localityDomainLocalizationAlloc(size, device, stream, 1); +} + +extern "C" TLLM_LOCALITY_DOMAIN_ALLOCATOR_EXPORT void trtllm_locality_domain0_free( + void* ptr, size_t size, int device, void* stream) noexcept +{ + torch_ext::localityDomainLocalizationFree(ptr, size, device, stream, 0); +} + +extern "C" TLLM_LOCALITY_DOMAIN_ALLOCATOR_EXPORT void trtllm_locality_domain1_free( + void* ptr, size_t size, int device, void* stream) noexcept +{ + torch_ext::localityDomainLocalizationFree(ptr, size, device, stream, 1); +} diff --git a/cpp/tests/unit_tests/runtime/CMakeLists.txt b/cpp/tests/unit_tests/runtime/CMakeLists.txt index 3a171ee39877..dd501abcce62 100644 --- a/cpp/tests/unit_tests/runtime/CMakeLists.txt +++ b/cpp/tests/unit_tests/runtime/CMakeLists.txt @@ -33,6 +33,8 @@ add_gtest(runtimeMpiUtilsTest mpiUtilsTest.cpp) add_gtest(runtimeKernelTest runtimeKernelTest.cpp) add_gtest(samplingConfigTest samplingConfigTest.cpp) add_gtest(samplingTest samplingTest.cpp) +add_gtest(localizationTest localizationTest.cu) +add_gtest(localityDomainPublicConfigTest localityDomainPublicConfigTest.cpp) add_gtest(sanitizerTest sanitizerTest.cpp) add_gtest(tllmBuffersTest tllmBuffersTest.cpp) add_gtest(transposeKVKernelTest transposeKVKernelTest.cpp) diff --git a/cpp/tests/unit_tests/runtime/localityDomainPublicConfigTest.cpp b/cpp/tests/unit_tests/runtime/localityDomainPublicConfigTest.cpp new file mode 100644 index 000000000000..3a6949d9e2da --- /dev/null +++ b/cpp/tests/unit_tests/runtime/localityDomainPublicConfigTest.cpp @@ -0,0 +1,93 @@ +/* + * 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 "tensorrt_llm/runtime/locality_domain/localityDomainResourceConfig.h" + +#include + +namespace tensorrt_llm::locality_domain::detail +{ + +#if CUDA_VERSION >= 13040 + +TEST(LocalityDomainPublicConfigTest, StrictSplitUsesDiscoveryWithoutCoschedulingOverride) +{ + SmResourceGroupParams const groupParams = makeStrictSmResourceGroupParams(); + for (int localityDomainId = 0; localityDomainId < kLocalityDomainCount; ++localityDomainId) + { + auto const& params = groupParams[localityDomainId]; + EXPECT_EQ(params.smCount, 0); + EXPECT_EQ(params.coscheduledSmCount, 0); + EXPECT_EQ(params.preferredCoscheduledSmCount, 0); + EXPECT_EQ(params.flags, CU_DEV_SM_RESOURCE_GROUP_LOCALITY_DOMAIN_ID); + EXPECT_EQ(params.localityDomainId, localityDomainId); + for (unsigned int const reserved : params.reserved) + { + EXPECT_EQ(reserved, 0); + } + } +} + +TEST(LocalityDomainPublicConfigTest, BalancedSplitUsesHalfDeviceWithBackfill) +{ + // Any total divisible by kSmCountAlignment * kLocalityDomainCount works here; the + // helper is pure arithmetic, so the value is deliberately not tied to a specific device. + constexpr unsigned int kTestSmCount = 128; + constexpr unsigned int kExpectedSmCount = 64; // 128 / 2, computed independently of the helper + SmResourceGroupParams const groupParams = makeBalancedSmResourceGroupParams(kTestSmCount); + for (int localityDomainId = 0; localityDomainId < kLocalityDomainCount; ++localityDomainId) + { + auto const& params = groupParams[localityDomainId]; + EXPECT_EQ(params.smCount, kExpectedSmCount); + EXPECT_EQ(params.coscheduledSmCount, 0); + EXPECT_EQ(params.preferredCoscheduledSmCount, 0); + EXPECT_EQ(params.flags, CU_DEV_SM_RESOURCE_GROUP_LOCALITY_DOMAIN_ID | CU_DEV_SM_RESOURCE_GROUP_BACKFILL); + EXPECT_EQ(params.localityDomainId, localityDomainId); + for (unsigned int const reserved : params.reserved) + { + EXPECT_EQ(reserved, 0); + } + } +} + +TEST(LocalityDomainPublicConfigTest, BalancedSplitRejectsOddPerGroupSmCount) +{ + EXPECT_TRUE(isBalancedSmCountValid(128)); // 64 per group, even + EXPECT_FALSE(isBalancedSmCountValid(126)); // 63 per group, odd + EXPECT_FALSE(isBalancedSmCountValid(2)); // below the per-group alignment minimum + EXPECT_FALSE(isBalancedSmCountValid(0)); +} + +TEST(LocalityDomainPublicConfigTest, StrictSplitAcceptsRemainderAndExactCover) +{ + EXPECT_TRUE(isStrictSplitCountValid(128, 60, 8)); // split with a remainder + EXPECT_TRUE(isStrictSplitCountValid(120, 60, 0)); // exact cover, no remainder + EXPECT_FALSE(isStrictSplitCountValid(128, 60, 6)); // remainder does not account for every SM + EXPECT_FALSE(isStrictSplitCountValid(118, 60, 0)); // per-group count exceeds half the device + EXPECT_FALSE(isStrictSplitCountValid(128, 0, 128)); // empty locality domain +} + +#else + +TEST(LocalityDomainPublicConfigTest, CUDAHeadersOlderThan134CompileWithoutPublicTypes) +{ + SUCCEED(); +} + +#endif + +} // namespace tensorrt_llm::locality_domain::detail diff --git a/cpp/tests/unit_tests/runtime/localizationTest.cu b/cpp/tests/unit_tests/runtime/localizationTest.cu new file mode 100644 index 000000000000..091f53b7adf7 --- /dev/null +++ b/cpp/tests/unit_tests/runtime/localizationTest.cu @@ -0,0 +1,522 @@ +/* + * 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 +#include +#include +#include +#include +#include +#include + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/logger.h" +#include "tensorrt_llm/runtime/locality_domain/locality_domain_utils.h" + +using namespace tensorrt_llm::locality_domain; + +// Memory allocation type for dual-stream tests +enum class MemoryAllocationType +{ + SAME, // Allocate on the same LOCALITY_DOMAIN as the stream + DIFFERENT, // Allocate on a different LOCALITY_DOMAIN from the stream + NORMAL // Use regular cudaMalloc (no LOCALITY_DOMAIN localization) +}; + +// CUDA kernel for int4 memory copy +// Each thread processes 8 int4 elements to hide memory latency +// Use strided access to maintain coalesced memory access within a warp +__global__ void memcpyInt4Kernel(int4* dst, int4 const* src, size_t numInt4Elements) +{ + size_t tid = blockIdx.x * blockDim.x + threadIdx.x; + size_t stride = blockDim.x * gridDim.x; + +// Process 8 int4 elements per thread with strided access +// This ensures threads in the same warp access consecutive memory addresses +#pragma unroll + for (int i = 0; i < 8; i++) + { + size_t idx = tid + i * stride; + if (idx < numInt4Elements) + { + dst[idx] = src[idx]; + } + } +} + +class LocalizationTest : public ::testing::Test +{ +protected: + static constexpr size_t kGiB = 1024ULL * 1024ULL * 1024ULL; + static constexpr size_t kMinPerfAllocationSize = 256ULL * 1024ULL * 1024ULL; + + struct DeviceAllocation + { + LocalizationHandle* handle = nullptr; + MemoryAllocationType type = MemoryAllocationType::NORMAL; + void* ptr = nullptr; + + DeviceAllocation() = default; + DeviceAllocation(DeviceAllocation const&) = delete; + DeviceAllocation& operator=(DeviceAllocation const&) = delete; + + ~DeviceAllocation() noexcept + { + resetNoThrow(); + } + + void allocate( + LocalizationHandle* newHandle, size_t size, MemoryAllocationType newType, int streamLocalityDomainId) + { + reset(); + handle = newHandle; + type = newType; + switch (type) + { + case MemoryAllocationType::SAME: handle->localityDomainMalloc(&ptr, size, streamLocalityDomainId); break; + case MemoryAllocationType::DIFFERENT: + handle->localityDomainMalloc(&ptr, size, 1 - streamLocalityDomainId); // 0->1, 1->0 + break; + case MemoryAllocationType::NORMAL: TLLM_CUDA_CHECK(cudaMalloc(&ptr, size)); break; + } + } + + void reset() + { + if (ptr == nullptr) + { + return; + } + if (type == MemoryAllocationType::NORMAL) + { + TLLM_CUDA_CHECK(cudaFree(ptr)); + } + else + { + handle->localityDomainFree(ptr); + } + ptr = nullptr; + } + + void resetNoThrow() noexcept + { + if (ptr == nullptr) + { + return; + } + if (type == MemoryAllocationType::NORMAL) + { + auto result = cudaFree(ptr); + if (result != cudaSuccess) + { + ADD_FAILURE() << "cudaFree failed during cleanup: " << cudaGetErrorString(result); + } + } + else + { + try + { + handle->localityDomainFree(ptr); + } + catch (...) + { + ADD_FAILURE() << "localityDomainFree failed during cleanup"; + } + } + ptr = nullptr; + } + }; + + struct StreamHolder + { + CUstream stream = nullptr; + + StreamHolder() = default; + StreamHolder(StreamHolder const&) = delete; + StreamHolder& operator=(StreamHolder const&) = delete; + }; + + std::unique_ptr mHandle; + bool mSupportsLocalityDomain = false; + + void SetUp() override + { + // Create handle + mHandle = std::make_unique(); + + // Check if LOCALITY_DOMAIN is supported + mSupportsLocalityDomain = mHandle->supportsLocalization(); + if (!mSupportsLocalityDomain) + { + TLLM_LOG_WARNING( + "LOCALITY_DOMAIN localization is not supported on this device, skipping LOCALITY_DOMAIN-specific " + "tests."); + } + } + + std::optional chooseAllocationSize(size_t requestedSize, int allocationCount) + { + size_t freeMem = 0; + size_t totalMem = 0; + TLLM_CUDA_CHECK(cudaMemGetInfo(&freeMem, &totalMem)); + (void) totalMem; + + size_t maxPerAllocation = freeMem / static_cast(allocationCount * 2); + size_t selectedSize = std::min(requestedSize, maxPerAllocation); + selectedSize = (selectedSize / sizeof(int4)) * sizeof(int4); + if (selectedSize < kMinPerfAllocationSize) + { + return std::nullopt; + } + return selectedSize; + } + + void TearDown() override + { + // Handle will be automatically destroyed + } + + // Helper function: Run memory copy test and return elapsed time (milliseconds) + float runMemcpyTest(void* dst, void const* src, size_t sizeBytes, CUstream stream, int iterations) + { + size_t numInt4Elements = sizeBytes / sizeof(int4); + int threadsPerBlock = 256; + // Each thread processes 8 int4 elements + int numBlocks = (numInt4Elements + threadsPerBlock * 8 - 1) / (threadsPerBlock * 8); + + cudaEvent_t start, stop; + TLLM_CUDA_CHECK(cudaEventCreate(&start)); + TLLM_CUDA_CHECK(cudaEventCreate(&stop)); + + // Warmup + memcpyInt4Kernel<<>>( + reinterpret_cast(dst), reinterpret_cast(src), numInt4Elements); + TLLM_CUDA_CHECK(cudaStreamSynchronize(stream)); + + // Actual performance test + TLLM_CUDA_CHECK(cudaEventRecord(start, stream)); + for (int i = 0; i < iterations; i++) + { + memcpyInt4Kernel<<>>( + reinterpret_cast(dst), reinterpret_cast(src), numInt4Elements); + } + TLLM_CUDA_CHECK(cudaEventRecord(stop, stream)); + TLLM_CUDA_CHECK(cudaEventSynchronize(stop)); + + float milliseconds = 0; + TLLM_CUDA_CHECK(cudaEventElapsedTime(&milliseconds, start, stop)); + + TLLM_CUDA_CHECK(cudaEventDestroy(start)); + TLLM_CUDA_CHECK(cudaEventDestroy(stop)); + + return milliseconds; + } + + // Helper function: Run dual-stream memory copy test and return elapsed time (milliseconds) + float runDualStreamMemcpyTest(void* dst0, void const* src0, void* dst1, void const* src1, size_t sizeBytes, + CUstream stream0, CUstream stream1, int iterations) + { + size_t numInt4Elements = sizeBytes / sizeof(int4); + int threadsPerBlock = 256; + // Each thread processes 8 int4 elements + int numBlocks = (numInt4Elements + threadsPerBlock * 8 - 1) / (threadsPerBlock * 8); + + cudaEvent_t start, stop, sync0, sync1; + TLLM_CUDA_CHECK(cudaEventCreate(&start)); + TLLM_CUDA_CHECK(cudaEventCreate(&stop)); + TLLM_CUDA_CHECK(cudaEventCreate(&sync0)); + TLLM_CUDA_CHECK(cudaEventCreate(&sync1)); + + // Warmup + memcpyInt4Kernel<<>>( + reinterpret_cast(dst0), reinterpret_cast(src0), numInt4Elements); + memcpyInt4Kernel<<>>( + reinterpret_cast(dst1), reinterpret_cast(src1), numInt4Elements); + TLLM_CUDA_CHECK(cudaStreamSynchronize(stream0)); + TLLM_CUDA_CHECK(cudaStreamSynchronize(stream1)); + + // Actual performance test + TLLM_CUDA_CHECK(cudaEventRecord(start, stream0)); // Record to default stream + + for (int i = 0; i < iterations; i++) + { + memcpyInt4Kernel<<>>( + reinterpret_cast(dst0), reinterpret_cast(src0), numInt4Elements); + memcpyInt4Kernel<<>>( + reinterpret_cast(dst1), reinterpret_cast(src1), numInt4Elements); + + // Record completion of both streams + TLLM_CUDA_CHECK(cudaEventRecord(sync0, stream0)); + TLLM_CUDA_CHECK(cudaEventRecord(sync1, stream1)); + + TLLM_CUDA_CHECK(cudaStreamWaitEvent(stream0, sync1)); + TLLM_CUDA_CHECK(cudaStreamWaitEvent(stream1, sync0)); + } + + TLLM_CUDA_CHECK(cudaEventRecord(stop, stream0)); + TLLM_CUDA_CHECK(cudaEventSynchronize(stop)); + + float milliseconds = 0; + TLLM_CUDA_CHECK(cudaEventElapsedTime(&milliseconds, start, stop)); + + TLLM_CUDA_CHECK(cudaEventDestroy(start)); + TLLM_CUDA_CHECK(cudaEventDestroy(stop)); + TLLM_CUDA_CHECK(cudaEventDestroy(sync0)); + TLLM_CUDA_CHECK(cudaEventDestroy(sync1)); + + return milliseconds; + } + + // Helper function: Run dual-stream performance test with configurable memory allocation + void runDualStreamPerformanceTest(MemoryAllocationType srcType, MemoryAllocationType dstType, char const* testName) + { + auto sizeBytesOpt = chooseAllocationSize(10ULL * kGiB, 4); + if (!sizeBytesOpt.has_value()) + { + GTEST_SKIP() << "Not enough free GPU memory for dual-stream performance test"; + } + size_t const sizeBytes = sizeBytesOpt.value(); + int const iterations = 10; + + // Allocate memory for stream 0 (LOCALITY_DOMAIN 0) + DeviceAllocation devSrc0; + DeviceAllocation devDst0; + devSrc0.allocate(mHandle.get(), sizeBytes, srcType, 0); + devDst0.allocate(mHandle.get(), sizeBytes, dstType, 0); + + // Allocate memory for stream 1 (LOCALITY_DOMAIN 1) + DeviceAllocation devSrc1; + DeviceAllocation devDst1; + devSrc1.allocate(mHandle.get(), sizeBytes, srcType, 1); + devDst1.allocate(mHandle.get(), sizeBytes, dstType, 1); + + // Initialize source memory + TLLM_CUDA_CHECK(cudaMemset(devSrc0.ptr, 0x42, sizeBytes)); + TLLM_CUDA_CHECK(cudaMemset(devSrc1.ptr, 0x43, sizeBytes)); + + // Create two localized streams on different locality domains + StreamHolder stream0{mHandle->createLocalizedStream(0)}; + StreamHolder stream1{mHandle->createLocalizedStream(1)}; + + // Run test + float totalTime = runDualStreamMemcpyTest( + devDst0.ptr, devSrc0.ptr, devDst1.ptr, devSrc1.ptr, sizeBytes, stream0.stream, stream1.stream, iterations); + float avgTime = totalTime / iterations; + // Bandwidth accounts for both read and write (2x data movement) on both streams. + float bandwidth = (2.0 * 2.0 * sizeBytes / static_cast(kGiB)) / (avgTime / 1000.0); // GB/s + + TLLM_LOG_INFO("%s:", testName); + TLLM_LOG_INFO(" Allocation size per buffer: %.2f GiB", sizeBytes / static_cast(kGiB)); + TLLM_LOG_INFO(" Total time: %.2f ms", totalTime); + TLLM_LOG_INFO(" Average time per iteration: %.2f ms", avgTime); + TLLM_LOG_INFO(" Bandwidth: %.2f GB/s", bandwidth); + } +}; + +// Test 1: Create and destroy handle +TEST_F(LocalizationTest, CreateAndDestroyHandle) +{ + // Handle is created in SetUp, just verify it's valid + ASSERT_NE(mHandle, nullptr); + + // TearDown will handle destruction automatically +} + +// Test 3: Allocate and free memory on LOCALITY_DOMAIN 0 +TEST_F(LocalizationTest, MallocFreeOnLocalityDomain0) +{ + if (!mSupportsLocalityDomain) + { + GTEST_SKIP() << "LOCALITY_DOMAIN not supported, skipping test"; + } + + void* devPtr = nullptr; + size_t size = 1024 * 1024; // 1 MB + + // Allocate memory on LOCALITY_DOMAIN 0 + ASSERT_NO_THROW(mHandle->localityDomainMalloc(&devPtr, size, 0)); + ASSERT_NE(devPtr, nullptr); + +#if CUDA_VERSION >= 13040 + int localityDomain = -1; + ASSERT_EQ(cuPointerGetAttribute( + &localityDomain, CU_POINTER_ATTRIBUTE_LOCALITY_DOMAIN_ORDINAL, reinterpret_cast(devPtr)), + CUDA_SUCCESS); + EXPECT_EQ(localityDomain, 0); +#endif + + // Free memory + ASSERT_NO_THROW(mHandle->localityDomainFree(devPtr)); +} + +// Test 4: Allocate and free memory on LOCALITY_DOMAIN 1 +TEST_F(LocalizationTest, MallocFreeOnLocalityDomain1) +{ + if (!mSupportsLocalityDomain) + { + GTEST_SKIP() << "LOCALITY_DOMAIN not supported, skipping test"; + } + + void* devPtr = nullptr; + size_t size = 1024 * 1024; // 1 MB + + // Allocate memory on LOCALITY_DOMAIN 1 + ASSERT_NO_THROW(mHandle->localityDomainMalloc(&devPtr, size, 1)); + ASSERT_NE(devPtr, nullptr); + +#if CUDA_VERSION >= 13040 + int localityDomain = -1; + ASSERT_EQ(cuPointerGetAttribute( + &localityDomain, CU_POINTER_ATTRIBUTE_LOCALITY_DOMAIN_ORDINAL, reinterpret_cast(devPtr)), + CUDA_SUCCESS); + EXPECT_EQ(localityDomain, 1); +#endif + + // Free memory + ASSERT_NO_THROW(mHandle->localityDomainFree(devPtr)); +} + +// Test 5: Create localized stream on LOCALITY_DOMAIN 0 +TEST_F(LocalizationTest, CreateLocalizedStreamLocalityDomain0) +{ + if (!mSupportsLocalityDomain) + { + GTEST_SKIP() << "LOCALITY_DOMAIN not supported, skipping test"; + } + + // Create stream localized to LOCALITY_DOMAIN 0 + CUstream stream = mHandle->createLocalizedStream(0); + ASSERT_NE(stream, nullptr); + EXPECT_EQ(mHandle->createLocalizedStream(0), stream); +} + +// Test 6: Create localized stream on LOCALITY_DOMAIN 1 +TEST_F(LocalizationTest, CreateLocalizedStreamLocalityDomain1) +{ + if (!mSupportsLocalityDomain) + { + GTEST_SKIP() << "LOCALITY_DOMAIN not supported, skipping test"; + } + + // Create stream localized to LOCALITY_DOMAIN 1 + CUstream stream = mHandle->createLocalizedStream(1); + ASSERT_NE(stream, nullptr); + EXPECT_EQ(mHandle->createLocalizedStream(1), stream); +} + +// 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. + +// Performance Test 1: Single stream, copy 20GB memory +TEST_F(LocalizationTest, DISABLED_PerformanceTestSingleStream20GB) +{ + auto sizeBytesOpt = chooseAllocationSize(20ULL * kGiB, 2); + if (!sizeBytesOpt.has_value()) + { + GTEST_SKIP() << "Not enough free GPU memory for single-stream performance test"; + } + size_t const sizeBytes = sizeBytesOpt.value(); + int const iterations = 10; + + // Allocate memory + DeviceAllocation devSrc; + DeviceAllocation devDst; + devSrc.allocate(mHandle.get(), sizeBytes, MemoryAllocationType::NORMAL, 0); + devDst.allocate(mHandle.get(), sizeBytes, MemoryAllocationType::NORMAL, 0); + + // Initialize source memory + TLLM_CUDA_CHECK(cudaMemset(devSrc.ptr, 0x42, sizeBytes)); + + // Create stream + StreamHolder stream; + TLLM_CUDA_CHECK(cudaStreamCreate(reinterpret_cast(&stream.stream))); + + // Run test + float totalTime = runMemcpyTest(devDst.ptr, devSrc.ptr, sizeBytes, stream.stream, iterations); + float avgTime = totalTime / iterations; + // Bandwidth accounts for both read and write (2x data movement) + float bandwidth = (2.0 * sizeBytes / static_cast(kGiB)) / (avgTime / 1000.0); // GB/s + + TLLM_LOG_INFO("Performance Test 1 - Single Stream (20GB):"); + TLLM_LOG_INFO(" Allocation size per buffer: %.2f GiB", sizeBytes / static_cast(kGiB)); + TLLM_LOG_INFO(" Total time: %.2f ms", totalTime); + TLLM_LOG_INFO(" Average time per iteration: %.2f ms", avgTime); + TLLM_LOG_INFO(" Bandwidth: %.2f GB/s", bandwidth); +} + +// Performance Test 2: Dual streams, source and destination on the same LOCALITY_DOMAIN as stream +TEST_F(LocalizationTest, DISABLED_PerformanceTestDualStreamSameLocalityDomain) +{ + if (!mSupportsLocalityDomain) + { + GTEST_SKIP() << "LOCALITY_DOMAIN not supported, skipping test"; + } + + runDualStreamPerformanceTest(MemoryAllocationType::SAME, MemoryAllocationType::SAME, + "Performance Test 2 - Dual Stream Same LOCALITY_DOMAIN (2x10GB)"); +} + +// Performance Test 3: Dual streams, source on same locality domains as stream and destination on different locality +// domains +TEST_F(LocalizationTest, DISABLED_PerformanceTestDualStreamSrcSameDstDifferentLocalityDomain) +{ + if (!mSupportsLocalityDomain) + { + GTEST_SKIP() << "LOCALITY_DOMAIN not supported, skipping test"; + } + + runDualStreamPerformanceTest(MemoryAllocationType::SAME, MemoryAllocationType::DIFFERENT, + "Performance Test 3 - Dual Stream Src Same Dst Different LOCALITY_DOMAIN (2x10GB)"); +} + +// Performance Test 3: Dual streams, source on different locality domains as stream and destination on same locality +// domains +TEST_F(LocalizationTest, DISABLED_PerformanceTestDualStreamSrcDifferentDstSameLocalityDomain) +{ + if (!mSupportsLocalityDomain) + { + GTEST_SKIP() << "LOCALITY_DOMAIN not supported, skipping test"; + } + + runDualStreamPerformanceTest(MemoryAllocationType::DIFFERENT, MemoryAllocationType::SAME, + "Performance Test 3 - Dual Stream Src Different Dst Same LOCALITY_DOMAIN (2x10GB)"); +} + +// Performance Test 3: Dual streams, source and destination on different locality domains as stream +TEST_F(LocalizationTest, DISABLED_PerformanceTestDualStreamDifferentLocalityDomain) +{ + if (!mSupportsLocalityDomain) + { + GTEST_SKIP() << "LOCALITY_DOMAIN not supported, skipping test"; + } + + runDualStreamPerformanceTest(MemoryAllocationType::DIFFERENT, MemoryAllocationType::DIFFERENT, + "Performance Test 3 - Dual Stream Src and Dst Different LOCALITY_DOMAIN (2x10GB)"); +} + +// Performance Test 4: Dual localized streams with regular cudaMalloc memory +TEST_F(LocalizationTest, DISABLED_PerformanceTestDualLocalizedStreamRegularMemory) +{ + if (!mSupportsLocalityDomain) + { + GTEST_SKIP() << "LOCALITY_DOMAIN not supported, skipping test"; + } + + runDualStreamPerformanceTest(MemoryAllocationType::NORMAL, MemoryAllocationType::NORMAL, + "Performance Test 4 - Dual Localized Stream with Regular Memory (2x10GB)"); +} From d650f25f80ee597489e84fe12525b7152ce66d40 Mon Sep 17 00:00:00 2001 From: Chulian Zhang <851104+zhangcl@users.noreply.github.com> Date: Sat, 22 Aug 2026 23:16:57 -0700 Subject: [PATCH 2/4] [None][fix] Address review comments on the locality domain runtime Drop the unrelated Green Context guard on remainder stream destruction, and give the test stream holder explicit ownership. Signed-off-by: Chulian Zhang <851104+zhangcl@users.noreply.github.com> --- .../locality_domain/locality_domain_utils.cpp | 2 +- cpp/tests/unit_tests/runtime/localizationTest.cu | 16 ++++++++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) 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 398871b82bcc..3306fbe1ba08 100644 --- a/cpp/tensorrt_llm/runtime/locality_domain/locality_domain_utils.cpp +++ b/cpp/tensorrt_llm/runtime/locality_domain/locality_domain_utils.cpp @@ -486,7 +486,7 @@ class GreenContextPartitions } } - if (mRemainderStream != nullptr && mApi.greenCtxDestroy != nullptr) + if (mRemainderStream != nullptr) { CUresult const result = cuStreamDestroy(mRemainderStream); if (result != CUDA_SUCCESS) diff --git a/cpp/tests/unit_tests/runtime/localizationTest.cu b/cpp/tests/unit_tests/runtime/localizationTest.cu index 091f53b7adf7..6acec1658c98 100644 --- a/cpp/tests/unit_tests/runtime/localizationTest.cu +++ b/cpp/tests/unit_tests/runtime/localizationTest.cu @@ -145,10 +145,25 @@ protected: struct StreamHolder { CUstream stream = nullptr; + //! True only for streams this holder created. Borrowed localized streams are owned by the + //! process-lifetime localization resource and must stay unowned here. + bool owned = false; StreamHolder() = default; StreamHolder(StreamHolder const&) = delete; StreamHolder& operator=(StreamHolder const&) = delete; + + ~StreamHolder() noexcept + { + if (owned && stream != nullptr) + { + auto const result = cudaStreamDestroy(reinterpret_cast(stream)); + if (result != cudaSuccess) + { + ADD_FAILURE() << "cudaStreamDestroy failed during cleanup: " << cudaGetErrorString(result); + } + } + } }; std::unique_ptr mHandle; @@ -445,6 +460,7 @@ TEST_F(LocalizationTest, DISABLED_PerformanceTestSingleStream20GB) // Create stream StreamHolder stream; TLLM_CUDA_CHECK(cudaStreamCreate(reinterpret_cast(&stream.stream))); + stream.owned = true; // Run test float totalTime = runMemcpyTest(devDst.ptr, devSrc.ptr, sizeBytes, stream.stream, iterations); From 128d35549183b6381adb584c0ceb5b5a325c76e9 Mon Sep 17 00:00:00 2001 From: Chulian Zhang <851104+zhangcl@users.noreply.github.com> Date: Tue, 25 Aug 2026 14:39:13 -0700 Subject: [PATCH 3/4] [None][fix] Do not link NVML into the locality domain build NVML is loaded with dlopen through NVMLWrapper, so linking it added a libnvidia-ml.so.1 dependency that broke nanobind stub generation on nodes without the driver library. The locality domain code calls no NVML. Signed-off-by: Chulian Zhang <851104+zhangcl@users.noreply.github.com> --- cpp/tensorrt_llm/runtime/CMakeLists.txt | 2 -- cpp/tensorrt_llm/thop/CMakeLists.txt | 3 +-- 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/cpp/tensorrt_llm/runtime/CMakeLists.txt b/cpp/tensorrt_llm/runtime/CMakeLists.txt index f7ca019a75dd..f09b7baff508 100644 --- a/cpp/tensorrt_llm/runtime/CMakeLists.txt +++ b/cpp/tensorrt_llm/runtime/CMakeLists.txt @@ -77,8 +77,6 @@ set_property(TARGET runtime_src PROPERTY POSITION_INDEPENDENT_CODE ON) set_property(TARGET runtime_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON) add_cuda_architectures(runtime_src 89) -target_link_libraries(runtime_src PUBLIC ${CUDA_NVML_LIB}) - target_include_directories(runtime_src PRIVATE ${MPI_C_INCLUDE_DIRS}) if(ENABLE_MULTI_DEVICE) diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index d31b3adbc1d6..718e95d4e28b 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -175,8 +175,7 @@ endif() if(ENABLE_MULTI_DEVICE) target_include_directories(th_common PUBLIC ${MPI_C_INCLUDE_DIRS}) - target_link_libraries(th_common PRIVATE ${MPI_C_LIBRARIES} ${NCCL_LIB} - CUDA::nvml) + target_link_libraries(th_common PRIVATE ${MPI_C_LIBRARIES} ${NCCL_LIB}) endif() if(NOT WIN32) From ec25f2540f0f2092ff01ed02df59dfe6906b41d2 Mon Sep 17 00:00:00 2001 From: Chulian Zhang <851104+zhangcl@users.noreply.github.com> Date: Wed, 26 Aug 2026 10:32:34 -0700 Subject: [PATCH 4/4] [None][feat] Add a side-effect-free locality domain capability query Constructing a LocalizationHandle creates a CUDA context and partitions the device, so it is unsuitable as a support probe. deviceSupportsLocalization() issues a driver attribute query only. Signed-off-by: Chulian Zhang <851104+zhangcl@users.noreply.github.com> --- .../nanobind/runtime/bindings.cpp | 6 +++++ .../locality_domain/locality_domain_utils.cpp | 25 +++++++++++++++++++ .../locality_domain/locality_domain_utils.h | 2 ++ 3 files changed, 33 insertions(+) diff --git a/cpp/tensorrt_llm/nanobind/runtime/bindings.cpp b/cpp/tensorrt_llm/nanobind/runtime/bindings.cpp index 85e38c555969..a624fea5baf1 100644 --- a/cpp/tensorrt_llm/nanobind/runtime/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/runtime/bindings.cpp @@ -448,6 +448,12 @@ void initBindings(nb::module_& m) "Get the borrowed process-lifetime remainder Green Context stream, or 0 when unavailable", nb::call_guard()); + m.def("device_supports_locality_domain", &tensorrt_llm::locality_domain::deviceSupportsLocalization, + nb::arg("device"), + "Return whether the device exposes public locality domains. Performs a driver attribute query only: it " + "creates no CUDA context and does not partition the device, so it is safe to call before selecting a device.", + nb::call_guard()); + // Initialize MoeLoadBalancer bindings initMoeBindings(m); // Initialize HostFunc bindings 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 3306fbe1ba08..011fe9bd5f02 100644 --- a/cpp/tensorrt_llm/runtime/locality_domain/locality_domain_utils.cpp +++ b/cpp/tensorrt_llm/runtime/locality_domain/locality_domain_utils.cpp @@ -1114,4 +1114,29 @@ CUstream LocalizationHandle::getReservedRemainderStream() return mImpl->getReservedRemainderStream(); } +bool deviceSupportsLocalization(int device) noexcept +{ +#if CUDA_VERSION >= 13040 + // cuInit() is idempotent and does not create a context. + if (cuInit(0) != CUDA_SUCCESS) + { + return false; + } + + CUdevice cuDevice{}; + if (cuDeviceGet(&cuDevice, device) != CUDA_SUCCESS) + { + return false; + } + + int localityDomainCount{}; + CUresult const result + = cuDeviceGetAttribute(&localityDomainCount, CU_DEVICE_ATTRIBUTE_LOCALITY_DOMAIN_COUNT, cuDevice); + return result == CUDA_SUCCESS && localityDomainCount >= kLocalityDomainCount; +#else + static_cast(device); + return false; +#endif +} + } // namespace tensorrt_llm::locality_domain 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 6ddf2e815fad..e57aa238d9f5 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,8 @@ class LocalizationHandle Localization* mImpl; }; +bool deviceSupportsLocalization(int device) noexcept; + } // namespace locality_domain } // namespace tensorrt_llm