From 09292b06a0d9e39ebd3f769f461ffdba6206d57d Mon Sep 17 00:00:00 2001 From: Morgan Funtowicz Date: Tue, 9 Jul 2024 12:15:41 +0000 Subject: [PATCH] updated logic and comment to detect cuda compute capabilities --- backends/trtllm/lib/backend.cpp | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/backends/trtllm/lib/backend.cpp b/backends/trtllm/lib/backend.cpp index 8e2cec82..1a3f598a 100644 --- a/backends/trtllm/lib/backend.cpp +++ b/backends/trtllm/lib/backend.cpp @@ -13,18 +13,17 @@ void huggingface::tgi::backends::InitializeBackend() { tle::ExecutorConfig huggingface::tgi::backends::GetExecutorConfig(const json &config, const std::string &workerPath) { tle::ExecutorConfig execConfig(1); - // TODO : Need to check for >= sm_80 (ampere) + // Get the compute capabilities of the current hardware nvmlDevice_t device; - int32_t cudaComputeCapabilitiesMajor, cudaComputeCapabilitiesMinor; - + int32_t cudaComputeCapabilitiesMajor = 0, cudaComputeCapabilitiesMinor = 0; if(nvmlDeviceGetHandleByIndex_v2(0, &device) == NVML_SUCCESS) { + SPDLOG_DEBUG("Successfully acquired nvmlDevice_t = 0"); if(nvmlDeviceGetCudaComputeCapability(device, &cudaComputeCapabilitiesMajor, &cudaComputeCapabilitiesMinor) == NVML_SUCCESS) { SPDLOG_INFO(FMT_STRING("Detected sm_{:d}{:d} compute capabilities"), cudaComputeCapabilitiesMajor, cudaComputeCapabilitiesMinor); - execConfig.setEnableChunkedContext(cudaComputeCapabilitiesMajor >= 8); } } - execConfig.setKvCacheConfig(tle::KvCacheConfig(true)); + // Single engine (TP = PP = 1) -> using leader mode (no MPI involved) if(config["/pretrained_config/mapping/world_size"_json_pointer].get() == 1){ SPDLOG_INFO("Detected single engine deployment, using leader mode"); execConfig.setParallelConfig(tle::ParallelConfig( @@ -34,7 +33,7 @@ tle::ExecutorConfig huggingface::tgi::backends::GetExecutorConfig(const json &co std::nullopt, std::nullopt )); - } else { + } else { // Multiple engines -> using orchestrator mode (MPI involved) SPDLOG_INFO("Detected sharded engine deployment, using orchestrator mode"); execConfig.setParallelConfig(tle::ParallelConfig( tle::CommunicationType::kMPI, @@ -44,6 +43,10 @@ tle::ExecutorConfig huggingface::tgi::backends::GetExecutorConfig(const json &co tle::OrchestratorConfig(true, workerPath) )); } + + // Define some configuration variables + execConfig.setKvCacheConfig(tle::KvCacheConfig(true)); + execConfig.setEnableChunkedContext(cudaComputeCapabilitiesMajor >= 8); return execConfig; }