From 262b004b139527d2243efa5bc28fc0f34f2a605e Mon Sep 17 00:00:00 2001 From: Tomer Shmilovich Date: Tue, 15 Jul 2025 00:39:15 -0700 Subject: [PATCH 1/7] Pass mode & directory Passing mode & directory parameters to relevant onboard & offload functions. Signed-off-by: Tomer Shmilovich --- .../batch_manager/kvCacheManager.h | 28 ++++++-- .../batch_manager/kvCacheTransferManager.h | 6 +- cpp/include/tensorrt_llm/executor/executor.h | 7 +- .../batch_manager/kvCacheManager.cpp | 66 ++++++++++++------- .../batch_manager/kvCacheTransferManager.cpp | 12 ++-- .../executor/kvCacheRetentionConfig.cpp | 4 +- cpp/tensorrt_llm/executor/serialization.cpp | 2 +- .../nanobind/executor/request.cpp | 4 +- cpp/tensorrt_llm/pybind/executor/request.cpp | 4 +- 9 files changed, 85 insertions(+), 48 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index 8940b160a153..fd2bcef3e4e6 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -445,6 +445,16 @@ class GenerationRequest return mKvCacheRetentionConfig.getDecodeDurationMs(); } + [[nodiscard]] executor::KvCacheTransferMode getTransferMode() const + { + return mKvCacheRetentionConfig.getTransferMode(); + } + + [[nodiscard]] std::string const& getDirectory() const + { + return mKvCacheRetentionConfig.getDirectory(); + } + // @brief Check whether the sequence uses cyclic KV cache. // @return `true` if we have begun overwriting the beginning of the sequence's KV cache. // @details If `true`, we cannot store the sequence's KV cache for reuse. @@ -702,11 +712,13 @@ class WindowBlockManager //! \brief Bring offloaded block from secondary to primary memory. //! \details Does nothing if block is already in primary memory. - void onboardBlock(BlockPtr const& offloadBlock); + void onboardBlock(BlockPtr const& offloadBlock, + executor::KvCacheTransferMode mode = executor::KvCacheTransferMode::DRAM, std::string const& directory = ""); //! \brief Bring block from primary to secondary memory. //! \details Does nothing if block is already in secondary memory. - void offloadBlock(BlockPtr const& block); + void offloadBlock(BlockPtr const& block, executor::KvCacheTransferMode mode = executor::KvCacheTransferMode::DRAM, + std::string const& directory = ""); //! \brief Find first new block that must be allocated for context phase and return it's concatenated token vectors. //! \details Only full blocks are considered. @@ -760,7 +772,8 @@ class WindowBlockManager //! \param sequence Sequence to which blocks are assigned. //! \return Number of matched tokens from loaded blocks. SizeType32 loadOrAllocateBlocks(std::vector const& blockKeys, SizeType32 numContextBlocks, - GenerationRequest& sequence, std::vector const& perBlockRetentions); + GenerationRequest& sequence, std::vector const& perBlockRetentions, + executor::KvCacheTransferMode mode = executor::KvCacheTransferMode::DRAM, std::string const& directory = ""); //! \brief Free block and all it's descendants. This makes block a claimed leaf block. void freeChildren(BlockPtr const& block, executor::RetentionPriority priority, @@ -769,7 +782,8 @@ class WindowBlockManager //! \brief Find block least likely to be reused, free it if necessary and return. [[nodiscard]] BlockPtr getFreeBlock( executor::RetentionPriority = executor::KvCacheRetentionConfig::kDefaultRetentionPriority, - std::optional durationMs = std::nullopt); + std::optional durationMs = std::nullopt, + executor::KvCacheTransferMode mode = executor::KvCacheTransferMode::DRAM, std::string const& directory = ""); //! \brief Free block from previous block and claim it from free blocks list. void claimLeafBlock(BlockPtr const& block, std::optional priority = std::nullopt, @@ -913,11 +927,13 @@ class BlockManager //! \brief Bring block from primary to secondary memory for window size. //! \details Does nothing if block is already in primary memory. - void onboardBlock(BlockPtr const& offloadBlock, SizeType32 windowSize); + void onboardBlock(BlockPtr const& offloadBlock, SizeType32 windowSize, + executor::KvCacheTransferMode mode = executor::KvCacheTransferMode::DRAM, std::string const& directory = ""); //! \brief Bring block from primary to secondary memory for window size. //! \details Does nothing if block is already in secondary memory. - void offloadBlock(BlockPtr const& block, SizeType32 windowSize); + void offloadBlock(BlockPtr const& block, SizeType32 windowSize, + executor::KvCacheTransferMode mode = executor::KvCacheTransferMode::DRAM, std::string const& directory = ""); void storeBlocks(std::vector const& blockKeys, std::vector const& blockIds, SizeType32 windowSize) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h index 2acfd6db60da..c2cfd49ecaf3 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h @@ -37,12 +37,12 @@ class KVCacheTransferManager //! \brief Onboard a block to gpu memory. void onboard(BlockPtr const& offloadBlock, BlockPtr const& block, std::vector const& pools, int numTokensToCopy = 0, executor::KvCacheTransferMode mode = executor::KvCacheTransferMode::DRAM, - std::optional directory = std::nullopt); + std::string const& directory = ""); //! \brief Offload a block to cpu memory. void offload(BlockPtr const& block, BlockPtr const& offloadBlock, std::vector const& pools, int numTokensToCopy = 0, executor::KvCacheTransferMode mode = executor::KvCacheTransferMode::DRAM, - std::optional directory = std::nullopt); + std::string const& directory = ""); //! \brief Synchronize the offload/onboard streams with the bufferManager stream. void syncTransfers(); @@ -67,7 +67,7 @@ class KVCacheTransferManager */ void copyBlock(BlockPtr const& src, BlockPtr const& dst, std::vector const& pools, bool isOffload, int numTokensToCopy = 0, executor::KvCacheTransferMode mode = executor::KvCacheTransferMode::DRAM, - std::optional directory = std::nullopt); + std::string const& directory = ""); runtime::BufferManager mBufferManager; runtime::BufferManager mOnboardManager; diff --git a/cpp/include/tensorrt_llm/executor/executor.h b/cpp/include/tensorrt_llm/executor/executor.h index 28c69074a3c2..8a4f4251d302 100644 --- a/cpp/include/tensorrt_llm/executor/executor.h +++ b/cpp/include/tensorrt_llm/executor/executor.h @@ -582,14 +582,13 @@ class KvCacheRetentionConfig explicit KvCacheRetentionConfig(std::vector const& tokenRangeRetentionPriorities, RetentionPriority decodeRetentionPriority = kDefaultRetentionPriority, std::optional decodeDurationMs = std::nullopt, - KvCacheTransferMode transferMode = KvCacheTransferMode::DRAM, - std::optional directory = std::nullopt); + KvCacheTransferMode transferMode = KvCacheTransferMode::DRAM, std::string const& directory = ""); [[nodiscard]] std::vector getTokenRangeRetentionConfigs() const; [[nodiscard]] RetentionPriority getDecodeRetentionPriority() const; [[nodiscard]] std::optional getDecodeDurationMs() const; [[nodiscard]] KvCacheTransferMode getTransferMode() const; - [[nodiscard]] std::optional getDirectory() const; + [[nodiscard]] std::string const& getDirectory() const; /// @brief Convert the token range data into an entry per kv block. Returns a tuple of vectors corresponding to the /// priorities and durations for each block. @@ -616,7 +615,7 @@ class KvCacheRetentionConfig /// @brief The transfer mode for the block. KvCacheTransferMode mTransferMode; /// @brief Name of the directory if transfer mode is GDS or POSIX_DEBUG_FALLBACK. - std::optional mDirectory; + std::string mDirectory; }; /// @brief A class that holds information about the request diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index 748fcbbe09dd..b347c35ec725 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -855,8 +855,9 @@ void WindowBlockManager::freeChildren( claimLeafBlock(block, priority, durationMs); } -BlockPtr WindowBlockManager::getFreeBlock( - executor::RetentionPriority priority, std::optional durationMs) +BlockPtr WindowBlockManager::getFreeBlock(executor::RetentionPriority priority, + std::optional durationMs, executor::KvCacheTransferMode mode, + std::string const& directory) { // eviction policy get free primary block auto [block, canOffload] = mEvictionPolicy->getFreeBlock(kPrimaryLevel); @@ -877,7 +878,7 @@ BlockPtr WindowBlockManager::getFreeBlock( mEvictionPolicy->claimBlock(block); // Offload block in primary memory before repurposing auto offloadBlock = std::get<0>(mEvictionPolicy->getFreeBlock(kSecondaryLevel)); - mTransferManager->offload(block, offloadBlock, mPools); + mTransferManager->offload(block, offloadBlock, mPools, 0, mode, directory); // swap linear block offsets (i.e. make block the offload block) block->swapMemoryPoolBlockOffset(offloadBlock); @@ -929,17 +930,20 @@ void BlockManager::setOffsets(tk::KVCacheIndex* offsetsPtr, nvinfer1::Dims const mWindowBlockManagers.at(windowSize).setOffsets(offsetsPtr, offsetsShape, beamIdx, blockIdx, blockId); } -void BlockManager::onboardBlock(BlockPtr const& offloadBlock, SizeType32 windowSize) +void BlockManager::onboardBlock(BlockPtr const& offloadBlock, SizeType32 windowSize, executor::KvCacheTransferMode mode, + std::string const& directory) { - mWindowBlockManagers.at(windowSize).onboardBlock(offloadBlock); + mWindowBlockManagers.at(windowSize).onboardBlock(offloadBlock, mode, directory); } -void WindowBlockManager::onboardBlock(BlockPtr const& offloadBlock) +void WindowBlockManager::onboardBlock( + BlockPtr const& offloadBlock, executor::KvCacheTransferMode mode, std::string const& directory) { if (mOnboardBlocks && !offloadBlock->isPrimary()) { - auto block = getFreeBlock(); - mTransferManager->onboard(offloadBlock, block, mPools); + auto block + = getFreeBlock(executor::KvCacheRetentionConfig::kDefaultRetentionPriority, std::nullopt, mode, directory); + mTransferManager->onboard(offloadBlock, block, mPools, 0, mode, directory); // swap linear block offsets (i.e. make block the offload block and vice versa) offloadBlock->swapMemoryPoolBlockOffset(block); @@ -954,12 +958,14 @@ void WindowBlockManager::onboardBlock(BlockPtr const& offloadBlock) } } -void BlockManager::offloadBlock(BlockPtr const& block, SizeType32 windowSize) +void BlockManager::offloadBlock( + BlockPtr const& block, SizeType32 windowSize, executor::KvCacheTransferMode mode, std::string const& directory) { - mWindowBlockManagers.at(windowSize).offloadBlock(block); + mWindowBlockManagers.at(windowSize).offloadBlock(block, mode, directory); } -void WindowBlockManager::offloadBlock(BlockPtr const& block) +void WindowBlockManager::offloadBlock( + BlockPtr const& block, executor::KvCacheTransferMode mode, std::string const& directory) { if (mOnboardBlocks && block->isPrimary()) { @@ -967,7 +973,7 @@ void WindowBlockManager::offloadBlock(BlockPtr const& block) auto offloadBlock = std::get<0>(mEvictionPolicy->getFreeBlock(kSecondaryLevel)); // If we're swapping a block to secondary memory, maintain the prior priority values. mEvictionPolicy->claimBlock(offloadBlock); - mTransferManager->offload(block, offloadBlock, mPools); + mTransferManager->offload(block, offloadBlock, mPools, 0, mode, directory); // swap linear block offsets (i.e. make block the offload block) block->swapMemoryPoolBlockOffset(offloadBlock); @@ -1021,7 +1027,8 @@ bool WindowBlockManager::blockInRadixTree(BlockPtr const& block) } SizeType32 WindowBlockManager::loadOrAllocateBlocks(std::vector const& blockKeys, SizeType32 numContextBlocks, - GenerationRequest& sequence, std::vector const& perBlockRetentions) + GenerationRequest& sequence, std::vector const& perBlockRetentions, + executor::KvCacheTransferMode mode, std::string const& directory) { SizeType32 numMatchedTokens{0}; auto searchRoot = mCachedBlocksRoot; @@ -1055,8 +1062,9 @@ SizeType32 WindowBlockManager::loadOrAllocateBlocks(std::vector const& if (matchingBlock->hasRefs() || !matchingBlock->isLeaf()) { // Somebody else is using block or it is not a leaf, copy reusable tokens - auto newBlock = getFreeBlock(matchingBlock->getPriority(), matchingBlock->getDurationMs()); - mTransferManager->onboard(matchingBlock, newBlock, mPools, numMatched); + auto newBlock + = getFreeBlock(matchingBlock->getPriority(), matchingBlock->getDurationMs(), mode, directory); + mTransferManager->onboard(matchingBlock, newBlock, mPools, numMatched, mode, directory); // TODO: (optional) Send out event matchingBlock = newBlock; TLLM_LOG_DEBUG("%s::loadOrAllocateBlocks - Copied partially filled block %d", mLogPrefix.c_str(), @@ -1080,7 +1088,7 @@ SizeType32 WindowBlockManager::loadOrAllocateBlocks(std::vector const& TLLM_LOG_DEBUG("%s::loadOrAllocateBlocks - Matched full block %d", mLogPrefix.c_str(), matchingBlockId); searchRoot = matchingBlock; } - onboardBlock(matchingBlock); + onboardBlock(matchingBlock, mode, directory); addBlockToAllBeams(matchingBlock, sequence); // TODO: only add once for reused blocks ++mReusedBlocks; @@ -1096,7 +1104,7 @@ SizeType32 WindowBlockManager::loadOrAllocateBlocks(std::vector const& // If we haven't set a priority, set it to the default priority level (low) auto freeBlock = getFreeBlock(perBlockRetentions[bi].retentionPriority.value_or( executor::KvCacheRetentionConfig::kDefaultRetentionPriority), - perBlockRetentions[bi].durationMs); + perBlockRetentions[bi].durationMs, mode, directory); addBlockToAllBeams(freeBlock, sequence); TLLM_LOG_DEBUG("%s::loadOrAllocateBlocks - No match, allocated new block %d for sequence %lu", mLogPrefix.c_str(), freeBlock->getBlockId(), sequence.getRequestId()); @@ -1122,7 +1130,7 @@ SizeType32 WindowBlockManager::loadOrAllocateBlocks(std::vector const& // If we haven't set a priority, set it to the default priority level (low) auto freeBlock = getFreeBlock(perBlockRetentions[bi].retentionPriority.value_or( executor::KvCacheRetentionConfig::kDefaultRetentionPriority), - perBlockRetentions[bi].durationMs); + perBlockRetentions[bi].durationMs, mode, directory); addBlockToBeam(freeBlock, sequence, beamIdx); if (blockItr != blockKeys.end()) { @@ -1195,9 +1203,20 @@ void WindowBlockManager::addSequence( auto perBlockRetentions = config.value_or(executor::KvCacheRetentionConfig()) .getPerBlockRetentionPriorityDuration(getTokensPerBlock(), inputLength); + auto mode = config.value_or(executor::KvCacheRetentionConfig()).getTransferMode(); + auto directory = config.value_or(executor::KvCacheRetentionConfig()).getDirectory(); + + if (mode != executor::KvCacheTransferMode::DRAM && directory.empty()) + { + TLLM_LOG_WARNING( + "Transfer mode %d specified without directory, falling back to DRAM mode", static_cast(mode)); + mode = executor::KvCacheTransferMode::DRAM; + } + TLLM_CHECK(perBlockRetentions.size() == (size_t) numContextBlocks); - auto const prepopulatedPromptLen = loadOrAllocateBlocks(blockKeys, numContextBlocks, sequence, perBlockRetentions); + auto const prepopulatedPromptLen + = loadOrAllocateBlocks(blockKeys, numContextBlocks, sequence, perBlockRetentions, mode, directory); mReusedTokens += static_cast(prepopulatedPromptLen); mTotalInputTokens += static_cast(uniqueTokens.size()); @@ -1287,7 +1306,8 @@ void WindowBlockManager::allocateBlock(GenerationRequest& sequence, bool shareAm if (shareAmongBeams) { // add same block to all beams - auto block = getFreeBlock(sequence.getDecodeRetentionPriority(), sequence.getDecodeDurationMs()); + auto block = getFreeBlock(sequence.getDecodeRetentionPriority(), sequence.getDecodeDurationMs(), + sequence.getTransferMode(), sequence.getDirectory()); for (auto beamIdx = 0; beamIdx < beamWidth; ++beamIdx) { addBlockToBeam(block, sequence, beamIdx); @@ -1298,7 +1318,8 @@ void WindowBlockManager::allocateBlock(GenerationRequest& sequence, bool shareAm // add different block to each beam for (auto beamIdx = 0; beamIdx < beamWidth; ++beamIdx) { - auto block = getFreeBlock(sequence.getDecodeRetentionPriority(), sequence.getDecodeDurationMs()); + auto block = getFreeBlock(sequence.getDecodeRetentionPriority(), sequence.getDecodeDurationMs(), + sequence.getTransferMode(), sequence.getDirectory()); addBlockToBeam(block, sequence, beamIdx); } } @@ -1399,7 +1420,8 @@ void WindowBlockManager::replaceSharedBlock(GenerationRequest& sequence, SizeTyp TLLM_CHECK_WITH_INFO(hasFreeBlocks(beamWidth), "Can't allocate new blocks. No free blocks left."); for (auto beamIdx = 0; beamIdx < beamWidth; ++beamIdx) { - auto block = getFreeBlock(); + auto block = getFreeBlock(executor::KvCacheRetentionConfig::kDefaultRetentionPriority, std::nullopt, + sequence.getTransferMode(), sequence.getDirectory()); block->incRefCount(); if (sequence.getCacheBlockIds(mWindowSize).at(beamIdx).size() == 0) { diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp index 4d49ebf9221b..c344b417328f 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp @@ -106,7 +106,7 @@ tr::ITensor::SharedPtr KVCacheTransferManager::computeBlockPointer( void KVCacheTransferManager::copyBlock(BlockPtr const& src, BlockPtr const& dst, std::vector const& pools, bool isOffload, int numTokensToCopy, executor::KvCacheTransferMode mode, - std::optional directory) + std::string const& directory) { TLLM_LOG_DEBUG("copyBlock entered: srcId=%d, dstId=%d, isOffload=%s, mode=%d", src->getBlockId(), dst->getBlockId(), (isOffload ? "true" : "false"), static_cast(mode)); @@ -165,13 +165,13 @@ void KVCacheTransferManager::copyBlock(BlockPtr const& src, BlockPtr const& dst, auto dstPtr = computeBlockPointer(dst, pools, poolIdx); TLLM_CHECK_WITH_INFO( - directory.has_value(), "Expected a directory path for KVCache offload, but none was provided."); + !directory.empty(), "Expected a directory path for KVCache offload, but none was provided."); int size = std::snprintf( - nullptr, 0, "%s/block_%d_pool_%zu.bin", directory.value().c_str(), src->getBlockId(), poolIdx); + nullptr, 0, "%s/block_%d_pool_%zu.bin", directory.c_str(), src->getBlockId(), poolIdx); std::string filename(size + 1, '\0'); - std::snprintf(filename.data(), filename.size(), "%s/block_%d_pool_%zu.bin", directory.value().c_str(), + std::snprintf(filename.data(), filename.size(), "%s/block_%d_pool_%zu.bin", directory.c_str(), src->getBlockId(), poolIdx); if (mode == executor::KvCacheTransferMode::POSIX_DEBUG_FALLBACK) @@ -272,7 +272,7 @@ void KVCacheTransferManager::copyBlock(BlockPtr const& src, BlockPtr const& dst, void KVCacheTransferManager::onboard(BlockPtr const& offloadBlock, BlockPtr const& block, std::vector const& pools, int numTokensToCopy, executor::KvCacheTransferMode mode, - std::optional directory) + std::string const& directory) { if (mode != executor::KvCacheTransferMode::DRAM && mPendingOffloads.find(offloadBlock->getBlockId()) == mPendingOffloads.end()) @@ -291,7 +291,7 @@ void KVCacheTransferManager::onboard(BlockPtr const& offloadBlock, BlockPtr cons void KVCacheTransferManager::offload(BlockPtr const& block, BlockPtr const& offloadBlock, std::vector const& pools, int numTokensToCopy, executor::KvCacheTransferMode mode, - std::optional directory) + std::string const& directory) { mPendingOffloads[block->getBlockId()] = tr::CudaEvent(); copyBlock(block, offloadBlock, pools, true, numTokensToCopy, mode, directory); diff --git a/cpp/tensorrt_llm/executor/kvCacheRetentionConfig.cpp b/cpp/tensorrt_llm/executor/kvCacheRetentionConfig.cpp index 82175f156016..1de31022391c 100644 --- a/cpp/tensorrt_llm/executor/kvCacheRetentionConfig.cpp +++ b/cpp/tensorrt_llm/executor/kvCacheRetentionConfig.cpp @@ -45,7 +45,7 @@ bool KvCacheRetentionConfig::TokenRangeRetentionConfig::operator==(TokenRangeRet KvCacheRetentionConfig::KvCacheRetentionConfig( std::vector const& tokenRangeRetentionPriorities, RetentionPriority decodeRetentionPriority, std::optional decodeDurationMs, - KvCacheTransferMode transferMode, std::optional directory) + KvCacheTransferMode transferMode, std::string const& directory) : mTokenRangeRetentionConfigs(std::vector(tokenRangeRetentionPriorities)) , mDecodeRetentionPriority{decodeRetentionPriority} , mDecodeDurationMs{decodeDurationMs} @@ -117,7 +117,7 @@ KvCacheTransferMode KvCacheRetentionConfig::getTransferMode() const return mTransferMode; } -std::optional KvCacheRetentionConfig::getDirectory() const +std::string const& KvCacheRetentionConfig::getDirectory() const { return mDirectory; } diff --git a/cpp/tensorrt_llm/executor/serialization.cpp b/cpp/tensorrt_llm/executor/serialization.cpp index bba8d19e2f6f..8763917802bf 100644 --- a/cpp/tensorrt_llm/executor/serialization.cpp +++ b/cpp/tensorrt_llm/executor/serialization.cpp @@ -1594,7 +1594,7 @@ KvCacheRetentionConfig Serialization::deserializeKvCacheRetentionConfig(std::ist auto decodePriority = su::deserialize(is); auto decodeDurationMs = intToDuration(su::deserialize>(is)); auto transferMode = su::deserialize(is); - auto directory = su::deserialize>(is); + auto directory = su::deserialize(is); return KvCacheRetentionConfig{ tokenRangeRetentionPriorities, decodePriority, decodeDurationMs, transferMode, directory}; diff --git a/cpp/tensorrt_llm/nanobind/executor/request.cpp b/cpp/tensorrt_llm/nanobind/executor/request.cpp index e56341b53e22..f9a7296f7a61 100644 --- a/cpp/tensorrt_llm/nanobind/executor/request.cpp +++ b/cpp/tensorrt_llm/nanobind/executor/request.cpp @@ -394,7 +394,7 @@ void initRequestBindings(nb::module_& m) new (&kvCacheRetentionConfig) tle::KvCacheRetentionConfig( nb::cast>(state[0]), nb::cast(state[1]), nb::cast>(state[2]), - nb::cast(state[3]), nb::cast>(state[4])); + nb::cast(state[3]), nb::cast(state[4])); }; auto kvCacheRetentionConfig = nb::class_(m, "KvCacheRetentionConfig"); @@ -417,7 +417,7 @@ void initRequestBindings(nb::module_& m) // TokenRangeRetentionPriority bindings have been defined. kvCacheRetentionConfig .def(nb::init, tle::RetentionPriority, - std::optional, tle::KvCacheTransferMode, std::optional>(), + std::optional, tle::KvCacheTransferMode, std::string>(), nb::arg("token_range_retention_configs"), nb::arg("decode_retention_priority") = tle::KvCacheRetentionConfig::kDefaultRetentionPriority, nb::arg("decode_duration_ms") = nb::none(), nb::arg("transfer_mode") = tle::KvCacheTransferMode::DRAM, diff --git a/cpp/tensorrt_llm/pybind/executor/request.cpp b/cpp/tensorrt_llm/pybind/executor/request.cpp index 904410c253b1..0a168e0ff0e5 100644 --- a/cpp/tensorrt_llm/pybind/executor/request.cpp +++ b/cpp/tensorrt_llm/pybind/executor/request.cpp @@ -364,7 +364,7 @@ void initRequestBindings(pybind11::module_& m) return tle::KvCacheRetentionConfig( state[0].cast>(), state[1].cast(), state[2].cast>(), - state[3].cast(), state[4].cast>()); + state[3].cast(), state[4].cast()); }; auto kvCacheRetentionConfig = py::class_(m, "KvCacheRetentionConfig"); @@ -386,7 +386,7 @@ void initRequestBindings(pybind11::module_& m) // TokenRangeRetentionPriority bindings have been defined. kvCacheRetentionConfig .def(py::init, tle::RetentionPriority, - std::optional, tle::KvCacheTransferMode, std::optional>(), + std::optional, tle::KvCacheTransferMode, std::string>(), py::arg("token_range_retention_configs"), py::arg("decode_retention_priority") = tle::KvCacheRetentionConfig::kDefaultRetentionPriority, py::arg("decode_duration_ms") = py::none(), From e47bdc9e03d11c19a6c945678122c52da5c5cd49 Mon Sep 17 00:00:00 2001 From: Tomer Shmilovich Date: Tue, 15 Jul 2025 00:40:30 -0700 Subject: [PATCH 2/7] Nixl Loopback Agent Implement class LoopbackAgent. Signed-off-by: Tomer Shmilovich --- .../batch_manager/kvCacheManager.h | 12 +- .../tensorrt_llm/executor/transferAgent.h | 99 +++++++++ .../batch_manager/kvCacheManager.cpp | 17 +- .../nixl_utils/CMakeLists.txt | 3 + .../nixl_utils/transferAgent.cpp | 194 ++++++++++++++++++ .../nixl_utils/transferAgent.h | 31 +++ 6 files changed, 350 insertions(+), 6 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index fd2bcef3e4e6..3ab694f69db5 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -22,6 +22,7 @@ #include "tensorrt_llm/batch_manager/llmRequest.h" // TODO forward declare #include "tensorrt_llm/common/optionalRef.h" #include "tensorrt_llm/executor/executor.h" +#include "tensorrt_llm/executor/transferAgent.h" #include "tensorrt_llm/kernels/kvCacheIndex.h" #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/common.h" @@ -42,6 +43,8 @@ #include #include +namespace kvc = tensorrt_llm::executor::kv_cache; + namespace tensorrt_llm::batch_manager::eviction_policy { class BaseEvictionPolicy; @@ -548,7 +551,8 @@ class WindowBlockManager SizeType32 blocksInSecondaryPool, SizeType32 maxNumSequences, std::shared_ptr stream, bool onboardBlocks, CacheType cacheType, std::optional secondaryOffloadMinPriority, std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, - std::shared_ptr kvCacheConnectorManager); + std::shared_ptr kvCacheConnectorManager, + std::shared_ptr loopbackAgent = nullptr); ~WindowBlockManager(); @@ -831,6 +835,8 @@ class WindowBlockManager std::shared_ptr mEvictionPolicy; // Event manager std::shared_ptr mEventManager; + // Pointer to parent loopback agent + std::shared_ptr mLoopbackAgent; // Transfer manager std::shared_ptr mTransferManager; @@ -878,7 +884,8 @@ class BlockManager std::optional secondaryOffloadMinPriority = std::nullopt, std::shared_ptr eventManager = nullptr, bool enablePartialReuse = true, bool copyOnPartialReuse = true, - std::shared_ptr kvCacheConnectorManager = nullptr); + std::shared_ptr kvCacheConnectorManager = nullptr, + std::optional agentConfig = std::nullopt); BlockManager(BlockManager const&) = delete; BlockManager& operator=(BlockManager const&) = delete; @@ -1172,6 +1179,7 @@ class BlockManager SizeType32 mNumLayers; SizeType32 mTokensPerBlock; std::shared_ptr mEventManager; + std::shared_ptr mLoopbackAgent; CudaStreamPtr mStream; CacheType mCacheType; diff --git a/cpp/include/tensorrt_llm/executor/transferAgent.h b/cpp/include/tensorrt_llm/executor/transferAgent.h index 0a2884b6f643..82b58eaeb0ad 100644 --- a/cpp/include/tensorrt_llm/executor/transferAgent.h +++ b/cpp/include/tensorrt_llm/executor/transferAgent.h @@ -17,10 +17,14 @@ #pragma once #include "tensorrt_llm/common/assert.h" +#include #include #include #include #include +#include +#include +#include #include #include @@ -109,6 +113,80 @@ class MemoryDescs std::vector mDescs; }; +class FileDesc +{ +public: + FileDesc(std::string const& filename, int flags, mode_t mode, size_t len) + : mLen{len} + { + int fd = ::open(filename.c_str(), flags, mode); + TLLM_CHECK_WITH_INFO(fd >= 0, "Failed to open '%s' (GDS)", filename.c_str()); + this->fd = fd; + } + + FileDesc(FileDesc&& other) noexcept + : fd(other.fd) + , mLen(other.mLen) + { + other.fd = -1; + other.mLen = 0; + } + + FileDesc& operator=(FileDesc&& other) noexcept + { + if (this != &other) + { + if (fd != -1) + ::close(fd); + fd = other.fd; + mLen = other.mLen; + other.fd = -1; + other.mLen = 0; + } + return *this; + } + + ~FileDesc() + { + if (fd != -1) + ::close(fd); + } + + [[nodiscard]] uint64_t getFd() const noexcept + { + return fd; + } + + [[nodiscard]] size_t getLen() const noexcept + { + return mLen; + } + + FileDesc(FileDesc const&) = delete; + FileDesc& operator=(FileDesc const&) = delete; + +private: + int fd; + size_t mLen; +}; + +class FileDescs +{ +public: + FileDescs(std::vector&& descs) + : mDescs(std::move(descs)) + { + } + + [[nodiscard]] std::vector const& getDescs() const noexcept + { + return mDescs; + } + +private: + std::vector mDescs; +}; + using TransferDescs = MemoryDescs; using RegisterDescs = MemoryDescs; using SyncMessage = std::string; @@ -221,6 +299,13 @@ class BaseTransferAgent virtual bool checkRemoteDescs(std::string const& name, MemoryDescs const& memoryDescs) = 0; }; +class BaseLoopbackAgent +{ +public: + virtual ~BaseLoopbackAgent() = default; + virtual void executeLoopbackRequest(MemoryDescs const& memoryDescs, FileDescs const& fileDescs, bool isOffload) = 0; +}; + class DynLibLoader final { public: @@ -264,4 +349,18 @@ template TLLM_THROW("Unknown backend name."); } +template +[[nodiscard]] std::shared_ptr makeLoopbackAgent(std::string const& backend, Args&&... args) +{ + if (backend == "nixl") + { + auto& loader = DynLibLoader::getInstance(); + using CreateNixlFuncType = std::shared_ptr (*)(BaseAgentConfig const*); + auto* func = loader.getFunctionPointer( + "libtensorrt_llm_nixl_wrapper.so", "createNixlLoopbackAgent"); + return func(std::forward(args)...); + } + TLLM_THROW("Unknown backend name."); +} + } // namespace tensorrt_llm::executor::kv_cache diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index b347c35ec725..037aef923f5e 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -41,6 +41,7 @@ namespace tc = tensorrt_llm::common; namespace tk = tensorrt_llm::kernels; namespace tle = tensorrt_llm::executor; +using namespace tle::kv_cache; using namespace tensorrt_llm::runtime; using namespace tensorrt_llm::batch_manager::kv_cache_manager; using namespace tensorrt_llm::batch_manager::eviction_policy; @@ -505,13 +506,19 @@ BlockManager::BlockManager(std::vector const& numKvHeadsPerLayer, Si SizeType32 sinkBubbleLength, bool onboardBlocks, CacheType cacheType, std::optional secondaryOffloadMinPriority, std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, - std::shared_ptr kvCacheConnectorManager) + std::shared_ptr kvCacheConnectorManager, + std::optional agentConfig) : mNumLayers{static_cast(numKvHeadsPerLayer.size())} , mTokensPerBlock{tokensPerBlock} , mEventManager{std::move(eventManager)} , mStream{stream} , mCacheType{cacheType} { + if (agentConfig.has_value()) + mLoopbackAgent = makeLoopbackAgent("nixl", &agentConfig.value()); + else + mLoopbackAgent = nullptr; + auto const uniqueWindowSizeToLayers = BaseKVCacheManager::groupLayersByWindowSize(maxAttentionWindowVec, mNumLayers); @@ -535,7 +542,7 @@ BlockManager::BlockManager(std::vector const& numKvHeadsPerLayer, Si mWindowBlockManagers.try_emplace(windowSize, dtype, windowSize, layersWithWindowSize, numKvHeadsPerLayer, sizePerHead, tokensPerBlock, allottedPrimaryBlocks, allottedSecondaryBlocks, maxNumSequences, stream, onboardBlocks, cacheType, secondaryOffloadMinPriority, mEventManager, enablePartialReuse, - copyOnPartialReuse, kvCacheConnectorManager); + copyOnPartialReuse, kvCacheConnectorManager, mLoopbackAgent); } auto const numAllPools = getNumPools(); @@ -578,7 +585,8 @@ WindowBlockManager::WindowBlockManager(nvinfer1::DataType dtype, SizeType32 wind SizeType32 maxNumSequences, std::shared_ptr stream, bool onboardBlocks, CacheType cacheType, std::optional secondaryOffloadMinPriority, std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, - std::shared_ptr kvCacheConnectorManager) + std::shared_ptr kvCacheConnectorManager, + std::shared_ptr loopbackAgent) : mDataType{dtype} , mWindowSize{windowSize} , mNumPrimaryBlocks{blocksInPrimaryPool} @@ -590,7 +598,8 @@ WindowBlockManager::WindowBlockManager(nvinfer1::DataType dtype, SizeType32 wind , mCachedBlocksRoot{std::make_shared(KVCacheBlock::kCachedBlocksRootId, tk::KVCacheIndex{0})} , mCacheType{cacheType} , mEventManager(std::move(eventManager)) - , mTransferManager{std::make_shared(mBufferManager)} + , mLoopbackAgent{loopbackAgent} + , mTransferManager{std::make_shared(mBufferManager, mLoopbackAgent)} , mAllocTotalBlocks{0} , mAllocNewBlocks{0} , mReusedBlocks{0} diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/CMakeLists.txt b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/CMakeLists.txt index de589ee325a6..78d95eac17ad 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/CMakeLists.txt +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/CMakeLists.txt @@ -37,4 +37,7 @@ if(NIXL_ROOT) # Link against all NIXL libraries target_link_libraries(${NIXL_WRAPPER_TARGET} PRIVATE NIXL::nixl) + # Link against CUDA + target_link_libraries(${NIXL_WRAPPER_TARGET} PRIVATE CUDA::cudart) + endif() diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp index 721cade13a54..252debaeb8d0 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp @@ -223,6 +223,16 @@ uint16_t getIncrmentPort(uint16_t basePort) return list; } +[[nodiscard]] nixl_reg_dlist_t NixlHelper::convertRegDlist(FileDescs const& descs) +{ + nixl_reg_dlist_t list(FILE_SEG); + for (auto const& desc : descs.getDescs()) + { + list.addDesc(nixlBlobDesc{0, desc.getLen(), desc.getFd()}); + } + return list; +} + [[nodiscard]] nixl_xfer_op_t NixlHelper::convert(TransferOp const& op) { switch (op) @@ -243,6 +253,62 @@ uint16_t getIncrmentPort(uint16_t basePort) return list; } +[[nodiscard]] nixl_xfer_dlist_t NixlHelper::convertXferDist(FileDescs const& descs) +{ + nixl_xfer_dlist_t list{FILE_SEG}; + for (auto const& desc : descs.getDescs()) + { + list.addDesc(nixlBasicDesc{0, desc.getLen(), desc.getFd()}); + } + return list; +} + +void NixlHelper::posixGpuToFileFallback(MemoryDescs const& memoryDescs, FileDescs const& fileDescs) +{ + auto const& memVec = memoryDescs.getDescs(); + auto const& fileVec = fileDescs.getDescs(); + std::size_t i; + + for (i = 0; i < std::min(memVec.size(), fileVec.size()); i++) + { + auto& memDesc = memVec[i]; + auto& fileDesc = fileVec[i]; + + ssize_t numBytes = static_cast(memDesc.getLen()); + std::vector hostBuffer(numBytes); + + cudaError_t cpyErr = cudaMemcpy( + hostBuffer.data(), reinterpret_cast(memDesc.getAddr()), numBytes, cudaMemcpyDeviceToHost); + TLLM_CHECK_WITH_INFO(cpyErr == cudaSuccess, "cudaMemcpy to host failed, error=%d", cpyErr); + + ssize_t written = ::write(fileDesc.getFd(), hostBuffer.data(), numBytes); + TLLM_CHECK_WITH_INFO(written >= 0, "POSIX write error=%zd", written); + } +} + +void NixlHelper::posixFileToGpuFallback(MemoryDescs const& memoryDescs, FileDescs const& fileDescs) +{ + auto const& memVec = memoryDescs.getDescs(); + auto const& fileVec = fileDescs.getDescs(); + std::size_t i; + + for (i = 0; i < std::min(memVec.size(), fileVec.size()); i++) + { + auto& memDesc = memVec[i]; + auto& fileDesc = fileVec[i]; + + ssize_t numBytes = static_cast(memDesc.getLen()); + std::vector hostBuffer(numBytes); + + ssize_t bytesRead = ::read(fileDesc.getFd(), hostBuffer.data(), numBytes); + TLLM_CHECK_WITH_INFO(bytesRead == numBytes, "POSIX read error=%zd", bytesRead); + + cudaError_t cpyErr = cudaMemcpy( + reinterpret_cast(memDesc.getAddr()), hostBuffer.data(), numBytes, cudaMemcpyHostToDevice); + TLLM_CHECK_WITH_INFO(cpyErr == cudaSuccess, "cudaMemcpy to device failed, error=%d", cpyErr); + } +} + NixlTransferStatus::NixlTransferStatus(nixlAgent* agent, nixlXferReqH* handle) : mRawAgent{agent} , mHandle{handle} @@ -457,6 +523,125 @@ NixlTransferAgent::~NixlTransferAgent() TLLM_LOG_DEBUG("NixlTransferAgent::~NixlTransferAgent"); } +NixlLoopbackAgent::NixlLoopbackAgent(BaseAgentConfig const& config) + : mName{config.mName} +{ + nixlAgentConfig nixlConfig{config.useProgThread}; + nixlBackendH* backend; + nixl_status_t status; + nixl_b_params_t init; + + mRawAgent = std::make_unique(config.mName, std::move(nixlConfig)); + init["batch_pool_size"] = std::to_string(8); + init["batch_limit"] = std::to_string(128); + init["max_request_size"] = std::to_string(16 * 1024 * 1024); + + status = mRawAgent->createBackend("GDS", init, backend); + if (status != NIXL_SUCCESS || !backend) + { + TLLM_THROW("Failed to create NIXL backend, status = %d", status); + } +} + +int NixlLoopbackAgent::registerMemory(MemoryDescs const& descs) +{ + nixl_status_t status = mRawAgent->registerMem(NixlHelper::convertRegDlist(descs)); + if (status != NIXL_SUCCESS) + return -1; + + return 0; +} + +int NixlLoopbackAgent::deregisterMemory(MemoryDescs const& descs) +{ + nixl_status_t status = mRawAgent->deregisterMem(NixlHelper::convertRegDlist(descs)); + if (status != NIXL_SUCCESS) + return -1; + + return 0; +} + +int NixlLoopbackAgent::registerFiles(FileDescs const& descs) +{ + nixl_status_t status = mRawAgent->registerMem(NixlHelper::convertRegDlist(descs)); + if (status != NIXL_SUCCESS) + return -1; + + return 0; +} + +int NixlLoopbackAgent::deregisterFiles(FileDescs const& descs) +{ + nixl_status_t status = mRawAgent->deregisterMem(NixlHelper::convertRegDlist(descs)); + if (status != NIXL_SUCCESS) + return -1; + + return 0; +} + +std::unique_ptr NixlLoopbackAgent::submitLoopbackRequests( + MemoryDescs const& memoryDescs, FileDescs const& fileDescs, bool isOffload) +{ + nixl_xfer_dlist_t vram_seg = NixlHelper::convertXferDist(memoryDescs); + nixl_xfer_dlist_t file_seg = NixlHelper::convertXferDist(fileDescs); + nixl_xfer_dlist_t& src = isOffload ? vram_seg : file_seg; + nixl_xfer_dlist_t& dst = isOffload ? file_seg : vram_seg; + nixl_xfer_op_t op = isOffload ? NIXL_WRITE : NIXL_READ; + nixlXferReqH* handle = nullptr; + + nixl_status_t status = mRawAgent->createXferReq(op, src, dst, mName, handle); + TLLM_CHECK(status == NIXL_SUCCESS && handle); + status = mRawAgent->postXferReq(handle); + TLLM_CHECK(status == NIXL_IN_PROG); + + return std::make_unique(mRawAgent.get(), handle); +} + +void NixlLoopbackAgent::executeLoopbackRequest( + MemoryDescs const& memoryDescs, FileDescs const& fileDescs, bool isOffload) +{ + bool fallback = false; + int ret; + + ret = this->registerFiles(fileDescs); + if (ret < 0) + { // register can fail if no GDS support + TLLM_LOG_DEBUG("NIXL GDS register files failed, using POSIX fallback"); + fallback = true; + } + else + { + ret = this->registerMemory(memoryDescs); + if (ret < 0) + { // register can fail if no GDS support + TLLM_LOG_DEBUG("NIXL GDS register memory failed, using POSIX fallback"); + this->deregisterFiles(fileDescs); + fallback = true; + } + } + + if (fallback) + { + if (isOffload) + { + NixlHelper::posixGpuToFileFallback(memoryDescs, fileDescs); + } + else + { + NixlHelper::posixFileToGpuFallback(memoryDescs, fileDescs); + } + + return; + } + + std::unique_ptr status = this->submitLoopbackRequests(memoryDescs, fileDescs, isOffload); + TLLM_CHECK_WITH_INFO(status != nullptr, "submitLoopbackRequests failed"); + status->wait(); + + this->deregisterMemory(memoryDescs); + this->deregisterFiles(fileDescs); +} + #if defined(__clang__) #pragma clang diagnostic push #pragma clang diagnostic ignored "-Wreturn-type-c-linkage" @@ -471,6 +656,15 @@ extern "C" } } +extern "C" +{ + std::shared_ptr createNixlLoopbackAgent(BaseAgentConfig const* config) + { + TLLM_CHECK(config); + return std::make_shared(*config); + } +} + #if defined(__clang__) #pragma clang diagnostic pop #endif diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h index 7288621e6d98..8d63ef656fc3 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h @@ -30,8 +30,12 @@ struct NixlHelper [[nodiscard]] static nixl_mem_t convert(MemoryType type); [[nodiscard]] static nixlBasicDesc convert(MemoryDesc const& desc); [[nodiscard]] static nixl_reg_dlist_t convertRegDlist(RegisterDescs const& descs); + [[nodiscard]] static nixl_reg_dlist_t convertRegDlist(FileDescs const& descs); [[nodiscard]] static nixl_xfer_op_t convert(TransferOp const& op); [[nodiscard]] static nixl_xfer_dlist_t convertXferDist(TransferDescs const& descs); + [[nodiscard]] static nixl_xfer_dlist_t convertXferDist(FileDescs const& descs); + static void posixGpuToFileFallback(MemoryDescs const& memoryDesc, FileDescs const& fileDescs); + static void posixFileToGpuFallback(MemoryDescs const& memoryDesc, FileDescs const& fileDescs); }; class NixlTransferStatus final : public TransferStatus @@ -97,6 +101,28 @@ class NixlTransferAgent final : public BaseTransferAgent std::vector mDRamDstBuffer; }; +class NixlLoopbackAgent final : public BaseLoopbackAgent +{ +public: + NixlLoopbackAgent(BaseAgentConfig const& config); + virtual ~NixlLoopbackAgent() = default; + + virtual void executeLoopbackRequest( + MemoryDescs const& memoryDescs, FileDescs const& fileDescs, bool isOffload) override; + +private: + int registerMemory(MemoryDescs const& descs); + int deregisterMemory(MemoryDescs const& descs); + int registerFiles(FileDescs const& descs); + int deregisterFiles(FileDescs const& descs); + + [[nodiscard]] std::unique_ptr submitLoopbackRequests( + MemoryDescs const& memoryDescs, FileDescs const& filedescs, bool isOffload); + + std::unique_ptr mRawAgent; + std::string mName; +}; + #if defined(__clang__) #pragma clang diagnostic push #pragma clang diagnostic ignored "-Wreturn-type-c-linkage" @@ -107,6 +133,11 @@ extern "C" [[nodiscard]] std::unique_ptr createNixlTransferAgent(BaseAgentConfig const* config); } +extern "C" +{ + [[nodiscard]] std::shared_ptr createNixlLoopbackAgent(BaseAgentConfig const* config); +} + #if defined(__clang__) #pragma clang diagnostic pop #endif From 8498b361fb645e2766d3b3cb70ab438007b8987f Mon Sep 17 00:00:00 2001 From: Tomer Shmilovich Date: Tue, 2 Sep 2025 08:37:52 -0700 Subject: [PATCH 3/7] copyBlock with NixlLoopbackAgent Signed-off-by: Tomer Shmilovich --- .../batch_manager/kvCacheTransferManager.h | 7 +- .../batch_manager/kvCacheTransferManager.cpp | 115 +++++------------- 2 files changed, 38 insertions(+), 84 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h index c2cfd49ecaf3..45f615cafe71 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h @@ -20,6 +20,7 @@ #include "tensorrt_llm/runtime/cudaEvent.h" namespace tr = tensorrt_llm::runtime; +namespace kvc = tensorrt_llm::executor::kv_cache; #pragma once @@ -32,7 +33,8 @@ namespace tensorrt_llm::batch_manager::kv_cache_manager class KVCacheTransferManager { public: - explicit KVCacheTransferManager(tr::BufferManager const& bufferManager); + explicit KVCacheTransferManager( + tr::BufferManager const& bufferManager, std::shared_ptr loopbackAgent = nullptr); //! \brief Onboard a block to gpu memory. void onboard(BlockPtr const& offloadBlock, BlockPtr const& block, std::vector const& pools, @@ -75,6 +77,9 @@ class KVCacheTransferManager // Track the block ids offloaded in this iteration. std::unordered_map mPendingOffloads; + // Reference to parent loopback agent + std::shared_ptr mLoopbackAgent; + int mDeviceId; }; } // namespace tensorrt_llm::batch_manager::kv_cache_manager diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp index c344b417328f..35868b35ac18 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp @@ -42,6 +42,7 @@ namespace tr = tensorrt_llm::runtime; namespace tk = tensorrt_llm::kernels; +namespace kvc = tensorrt_llm::executor::kv_cache; namespace tensorrt_llm::batch_manager::kv_cache_manager { @@ -86,11 +87,15 @@ static bool fileToGpuPosix(tr::ITensor::SharedPtr const& dstPtr, std::string con return true; } -KVCacheTransferManager::KVCacheTransferManager(tr::BufferManager const& bufferManager) +KVCacheTransferManager::KVCacheTransferManager( + tr::BufferManager const& bufferManager, std::shared_ptr loopbackAgent) : mBufferManager{bufferManager} , mOnboardManager(std::make_shared()) , mOffloadManager(std::make_shared()) + , mLoopbackAgent{loopbackAgent} { + TLLM_CUDA_CHECK(cudaGetDevice(&mDeviceId)); + TLLM_CHECK(mDeviceId != -1); } tr::ITensor::SharedPtr KVCacheTransferManager::computeBlockPointer( @@ -159,114 +164,58 @@ void KVCacheTransferManager::copyBlock(BlockPtr const& src, BlockPtr const& dst, return; } + std::vector fileBlobs; + std::vector memoryBlobs; + for (size_t poolIdx = 0; poolIdx < pools.size(); ++poolIdx) { - auto srcPtr = computeBlockPointer(src, pools, poolIdx); - auto dstPtr = computeBlockPointer(dst, pools, poolIdx); + auto ptr = isOffload ? computeBlockPointer(src, pools, poolIdx) : computeBlockPointer(dst, pools, poolIdx); + auto block_id = src->getBlockId(); TLLM_CHECK_WITH_INFO( !directory.empty(), "Expected a directory path for KVCache offload, but none was provided."); - int size = std::snprintf( - nullptr, 0, "%s/block_%d_pool_%zu.bin", directory.c_str(), src->getBlockId(), poolIdx); - - std::string filename(size + 1, '\0'); - std::snprintf(filename.data(), filename.size(), "%s/block_%d_pool_%zu.bin", directory.c_str(), - src->getBlockId(), poolIdx); + int size = std::snprintf(nullptr, 0, "%s/block_%d_pool_%zu.bin", directory.c_str(), block_id, poolIdx); + std::string filename; + filename.resize(size + 1); + std::snprintf( + filename.data(), filename.size(), "%s/block_%d_pool_%zu.bin", directory.c_str(), block_id, poolIdx); if (mode == executor::KvCacheTransferMode::POSIX_DEBUG_FALLBACK) { TLLM_LOG_INFO("Forcing POSIX fallback for file: %s", filename.c_str()); if (isOffload) { - gpuToFilePosix(srcPtr, filename); + gpuToFilePosix(ptr, filename); } else { - fileToGpuPosix(dstPtr, filename); + fileToGpuPosix(ptr, filename); } continue; } - - int openFlags = isOffload ? (O_CREAT | O_WRONLY) : O_RDONLY; - int fd = ::open(filename.c_str(), openFlags, 0664); - if (fd < 0) + else if (mode == executor::KvCacheTransferMode::GDS) { - TLLM_LOG_ERROR( - "Failed to open '%s' for %s; fallback POSIX", filename.c_str(), (isOffload ? "writing" : "reading")); - if (isOffload) - { - gpuToFilePosix(srcPtr, filename); - } - else - { - fileToGpuPosix(dstPtr, filename); - } - continue; + int openFlags = isOffload ? (O_CREAT | O_WRONLY) : O_RDONLY; + fileBlobs.emplace_back(filename, openFlags, 0664, ptr->getSizeInBytes()); + memoryBlobs.emplace_back(ptr->data(), ptr->getSizeInBytes(), mDeviceId); } + } -#ifdef ENABLE_CUFILE - CUfileDescr_t cufileDesc = {}; - cufileDesc.type = CU_FILE_HANDLE_TYPE_OPAQUE_FD; - cufileDesc.handle.fd = fd; - - CUfileHandle_t cufileHandle; - CUfileError_t status = cuFileHandleRegister(&cufileHandle, &cufileDesc); - if (status.err != CU_FILE_SUCCESS) + if (mode == executor::KvCacheTransferMode::GDS) + { + if (mLoopbackAgent == nullptr) { - // Fallback to POSIX - TLLM_LOG_WARN( - "cuFileHandleRegister failed (err=%d). Falling back to POSIX for '%s'", status.err, filename.c_str()); - ::close(fd); - if (isOffload) - gpuToFilePosix(srcPtr, filename); - else - fileToGpuPosix(dstPtr, filename); - continue; + TLLM_LOG_DEBUG("KVCacheTransferManager: creating mLoopbackAgent lazily"); + kvc::BaseAgentConfig config{std::string("GDSAgent"), true, true}; + mLoopbackAgent = kvc::makeLoopbackAgent("nixl", &config); } - ssize_t numBytes = static_cast(srcPtr->getSizeInBytes()); - if (isOffload) - { - ssize_t written = cuFileWrite(cufileHandle, srcPtr->data(), numBytes, 0, 0); - if (written < 0) - { - TLLM_LOG_ERROR("cuFileWrite error=%zd. Fallback to POSIX", written); - cuFileHandleDeregister(cufileHandle); - ::close(fd); - gpuToFilePosix(srcPtr, filename); - continue; - } - } - else - { - ssize_t readCount = cuFileRead(cufileHandle, dstPtr->data(), numBytes, 0, 0); - if (readCount < 0) - { - TLLM_LOG_ERROR("cuFileRead error=%zd. Fallback to POSIX", readCount); - cuFileHandleDeregister(cufileHandle); - ::close(fd); - fileToGpuPosix(dstPtr, filename); - continue; - } - } + kvc::FileDescs fileDescs(std::move(fileBlobs)); + kvc::MemoryDescs memoryDescs(kvc::MemoryType::kVRAM, memoryBlobs); - cuFileHandleDeregister(cufileHandle); - ::close(fd); -#else - // If GDS isn't enabled, fallback to POSIX automatically - TLLM_LOG_DEBUG("ENABLE_CUFILE=OFF, so fallback to POSIX for %s", filename.c_str()); - ::close(fd); // close the file opened for GDS - if (isOffload) - { - gpuToFilePosix(srcPtr, filename); - } - else - { - fileToGpuPosix(dstPtr, filename); - } -#endif + mLoopbackAgent->executeLoopbackRequest(memoryDescs, fileDescs, isOffload); } } From 2cc007219bb1f8d396636baae399cf6a9e44889a Mon Sep 17 00:00:00 2001 From: Tomer Shmilovich Date: Tue, 15 Jul 2025 00:39:16 -0700 Subject: [PATCH 4/7] GDS_MT backend support for LoopbackAgent Signed-off-by: Tomer Shmilovich --- cpp/include/tensorrt_llm/executor/transferAgent.h | 1 + .../cache_transmission/nixl_utils/transferAgent.cpp | 13 ++++++++++--- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/cpp/include/tensorrt_llm/executor/transferAgent.h b/cpp/include/tensorrt_llm/executor/transferAgent.h index 82b58eaeb0ad..c56278674681 100644 --- a/cpp/include/tensorrt_llm/executor/transferAgent.h +++ b/cpp/include/tensorrt_llm/executor/transferAgent.h @@ -273,6 +273,7 @@ struct BaseAgentConfig { std::string mName; bool useProgThread; + bool multiThread; }; class BaseTransferAgent diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp index 252debaeb8d0..4aae429374f2 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp @@ -536,10 +536,17 @@ NixlLoopbackAgent::NixlLoopbackAgent(BaseAgentConfig const& config) init["batch_limit"] = std::to_string(128); init["max_request_size"] = std::to_string(16 * 1024 * 1024); - status = mRawAgent->createBackend("GDS", init, backend); - if (status != NIXL_SUCCESS || !backend) + if (config.multiThread) { - TLLM_THROW("Failed to create NIXL backend, status = %d", status); + status = mRawAgent->createBackend("GDS_MT", init, backend); + if (status != NIXL_SUCCESS || !backend) + TLLM_THROW("Failed to create NIXL GDS_MT backend, status = %d", status); + } + else + { + status = mRawAgent->createBackend("GDS", init, backend); + if (status != NIXL_SUCCESS || !backend) + TLLM_THROW("Failed to create NIXL GDS backend, status = %d", status); } } From 283263a3a96bb535dbfb55c63e6b8879efd85836 Mon Sep 17 00:00:00 2001 From: Guy Lev Date: Thu, 7 Aug 2025 14:54:10 +0300 Subject: [PATCH 5/7] Add GDS memory mode tests to kvCacheManagerTest Signed-off-by: Guy Lev --- .../batch_manager/kvCacheManagerTest.cpp | 94 ++++++++++++++++--- 1 file changed, 79 insertions(+), 15 deletions(-) diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp index 0a52ae84852e..eda0638bb51d 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp @@ -18,6 +18,7 @@ #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/memoryUtils.h" +#include "tensorrt_llm/executor/transferAgent.h" #include "tensorrt_llm/executor/types.h" #include "tensorrt_llm/kernels/kvCacheIndex.h" #include "tensorrt_llm/kernels/kvCacheUtils.h" @@ -32,6 +33,7 @@ #include #include #include +#include #include #include #include @@ -45,6 +47,7 @@ namespace tk = tensorrt_llm::kernels; namespace tlk = tensorrt_llm::batch_manager::kv_cache_manager; namespace tle = tensorrt_llm::executor; namespace tr = tensorrt_llm::runtime; +namespace fs = std::filesystem; using BlocksPerWindow = std::map>; @@ -182,7 +185,39 @@ TEST_F(KVCacheManagerTest, BlockManagerTest) std::runtime_error); } -template +template +void writePatternToOffloadedBlocksDRAM(T* rawBlockPtr, int blockSize, int mask) +{ + for (int i = 0; i < blockSize; ++i) + { + rawBlockPtr[i] = i & mask; + } +} + +template +void writePatternToOffloadedBlocksGDS( + std::string const& directory, int blockId, SizeType32 numPools, int blockSize, int mask) +{ + for (size_t poolIdx = 0; poolIdx < numPools; ++poolIdx) + { + std::string filename + = directory + "/block_" + std::to_string(blockId) + "_pool_" + std::to_string(poolIdx) + ".bin"; + int fd = ::open(filename.c_str(), O_WRONLY); + if (fd >= 0) + { + auto poolBlockSize = blockSize / numPools; + std::vector buffer(poolBlockSize); + for (int i = 0; i < poolBlockSize; ++i) + { + buffer[i] = i & mask; + } + ::write(fd, buffer.data(), poolBlockSize * sizeof(T)); + ::close(fd); + } + } +} + +template void runPartialCopyTest() { auto constexpr numLayers = 12; @@ -202,6 +237,16 @@ void runPartialCopyTest() auto constexpr maxAttentionWindowAllLayer = 4096; auto constexpr sinkTokenLen = 0; auto constexpr canUseOneMoreBlock = true; + std::string directory = ""; + static int file_num = 0; + + if constexpr (transferMode == KvCacheTransferMode::GDS) + { + std::string filename = std::string("test_copy") + std::to_string(file_num++); + auto dirPath = fs::absolute(filename); + fs::create_directories(dirPath); + directory = dirPath.string(); + } SizeType32 constexpr maxNewTokens{0}; auto constexpr beamWidth = 1; @@ -256,7 +301,7 @@ void runPartialCopyTest() auto block = blockManager.getBlockById(cacheBlockId, maxAttentionWindow); EXPECT_TRUE(block->isPrimary()); // offload so we can write to block in CPU code - blockManager.offloadBlock(block, maxAttentionWindow); + blockManager.offloadBlock(block, maxAttentionWindow, transferMode, directory); EXPECT_FALSE(block->isPrimary()); // need to sync so D2H transfer is done before accessing blocks EXPECT_EQ(cudaDeviceSynchronize(), cudaSuccess); @@ -264,12 +309,19 @@ void runPartialCopyTest() auto memoryPoolIndex = block->getMemoryPoolBlockIndex(); auto blockPtr{tr::ITensor::slice(secondaryPoolPtr, memoryPoolIndex, 1)}; auto rawBlockPtr = reinterpret_cast(blockPtr->data()); - for (int i = 0; i < blockSize; ++i) + // Write value + if constexpr (transferMode == KvCacheTransferMode::DRAM) { - rawBlockPtr[i] = i & mask; + writePatternToOffloadedBlocksDRAM(rawBlockPtr, blockSize, mask); + } + else if constexpr (transferMode == KvCacheTransferMode::GDS) + { + auto block_id = block->getBlockId(); + auto numPools = blockManager.getNumPools(false); + writePatternToOffloadedBlocksGDS(directory, block_id, numPools, blockSize, mask); } // onboard - blockManager.onboardBlock(block, maxAttentionWindow); + blockManager.onboardBlock(block, maxAttentionWindow, transferMode, directory); EXPECT_TRUE(block->isPrimary()); EXPECT_EQ(cudaDeviceSynchronize(), cudaSuccess); EXPECT_TRUE(blockManager.verifyQueueIntegrity(maxAttentionWindow)); @@ -344,60 +396,72 @@ void runPartialCopyTest() } } EXPECT_EQ(numBad, 0); - blockManager.onboardBlock(block2, maxAttentionWindow); + blockManager.onboardBlock(block2, maxAttentionWindow, transferMode, directory); EXPECT_TRUE(block2->isPrimary()); EXPECT_EQ(cudaDeviceSynchronize(), cudaSuccess); blockManager.releaseBlocks(seq1, llmRequest1); blockManager.releaseBlocks(seq2, llmRequest2); + + if constexpr (transferMode == KvCacheTransferMode::GDS) + fs::remove_all(directory); } TEST_F(KVCacheManagerTest, BlockManagerTestPartialCopyINT64) { - runPartialCopyTest(); + runPartialCopyTest(); + runPartialCopyTest(); } TEST_F(KVCacheManagerTest, BlockManagerTestPartialCopyINT32) { - runPartialCopyTest(); + runPartialCopyTest(); + runPartialCopyTest(); } TEST_F(KVCacheManagerTest, BlockManagerTestPartialCopyFLOAT) { - runPartialCopyTest(); + runPartialCopyTest(); + runPartialCopyTest(); } #ifdef ENABLE_BF16 TEST_F(KVCacheManagerTest, BlockManagerTestPartialCopyBF16) { - runPartialCopyTest(); + runPartialCopyTest(); + runPartialCopyTest(); } #endif TEST_F(KVCacheManagerTest, BlockManagerTestPartialCopyHALF) { - runPartialCopyTest(); + runPartialCopyTest(); + runPartialCopyTest(); } TEST_F(KVCacheManagerTest, BlockManagerTestPartialCopyBOOL) { - runPartialCopyTest(); + runPartialCopyTest(); + runPartialCopyTest(); } TEST_F(KVCacheManagerTest, BlockManagerTestPartialCopyUINT8) { - runPartialCopyTest(); + runPartialCopyTest(); + runPartialCopyTest(); } TEST_F(KVCacheManagerTest, BlockManagerTestPartialCopyINT8) { - runPartialCopyTest(); + runPartialCopyTest(); + runPartialCopyTest(); } #ifdef ENABLE_FP8 TEST_F(KVCacheManagerTest, BlockManagerTestPartialCopyFP8) { - runPartialCopyTest(); + runPartialCopyTest(); + runPartialCopyTest(); } #endif From e0d7a83aaff69c6694a4809734d5727badc138d9 Mon Sep 17 00:00:00 2001 From: Guy Lev Date: Mon, 11 Aug 2025 18:17:30 +0300 Subject: [PATCH 6/7] Add LoopbackAgent tests to transferAgentTest Signed-off-by: Guy Lev --- .../unit_tests/executor/transferAgentTest.cpp | 119 ++++++++++++++++++ 1 file changed, 119 insertions(+) diff --git a/cpp/tests/unit_tests/executor/transferAgentTest.cpp b/cpp/tests/unit_tests/executor/transferAgentTest.cpp index 4745e8e40b12..0ffa84e7a5b3 100644 --- a/cpp/tests/unit_tests/executor/transferAgentTest.cpp +++ b/cpp/tests/unit_tests/executor/transferAgentTest.cpp @@ -16,6 +16,10 @@ #include #include +#include + +namespace fs = std::filesystem; + using namespace tensorrt_llm::executor::kv_cache; class RegisteredHostMemory @@ -341,3 +345,118 @@ TEST_F(TransferAgentTest, SyncMessage) nixlAgent0->invalidateRemoteAgent(agent1); nixlAgent1->invalidateRemoteAgent(agent0); } + +class LoopbackAgentTest : public ::testing::Test, + public ::testing::WithParamInterface // NOLINT(cppcoreguidelines-pro-type-member-init) +{ +public: + void SetUp() override + { + static int file_num = 0; + std::string filename = std::string("test_agent") + std::to_string(file_num++); + auto dirPath = fs::absolute(filename); + std::error_code ec; + fs::create_directories(dirPath, ec); + TLLM_CHECK_WITH_INFO(!ec, "Failed to create test directory: %s", ec.message().c_str()); + mDirectory = dirPath.string(); + } + + void TearDown() override + { + std::error_code ec; + fs::remove_all(mDirectory, ec); + if (ec) + std::cerr << "Warning: Failed to clean up test directory: " << ec.message() << std::endl; + } + + [[nodiscard]] std::shared_ptr makeLoopbackAgent(BaseAgentConfig const& config) + { + return tensorrt_llm::executor::kv_cache::makeLoopbackAgent("nixl", &config); + } + + [[nodiscard]] std::string getDirectory() const + { + return mDirectory; + } + +private: + std::string mDirectory; +}; + +TEST_P(LoopbackAgentTest, FileToGpu) +{ + std::string const agentName{"loopbackAgent"}; + BaseAgentConfig config{agentName, true, GetParam()}; + auto loopbackAgent = makeLoopbackAgent(config); + + TLLM_CHECK(loopbackAgent); + + std::vector memory(100, 1); + char* cuda_mem; + TLLM_CUDA_CHECK(cudaMalloc(&cuda_mem, 100)); + TLLM_CUDA_CHECK(cudaMemcpy(cuda_mem, memory.data(), 100, cudaMemcpyHostToDevice)); + std::string filename = getDirectory() + std::string("/file2gpu.bin"); + + int fd = ::open(filename.c_str(), O_CREAT | O_WRONLY, 0664); + TLLM_CHECK_WITH_INFO(fd >= 0, "Failed to open '%s' for writing", filename.c_str()); + + std::vector fileData(100, 10); + ssize_t bytesWritten = ::write(fd, fileData.data(), fileData.size()); + TLLM_CHECK_WITH_INFO(bytesWritten == static_cast(fileData.size()), "Failed to write to file"); + ::close(fd); + + { + MemoryDesc mem_desc(cuda_mem, 100, 0); + MemoryDescs memDescs{MemoryType::kVRAM, {mem_desc}}; + + std::vector fileDescVec; + fileDescVec.emplace_back(filename, O_RDONLY, 0664, 100); + FileDescs fileDescs{std::move(fileDescVec)}; + + loopbackAgent->executeLoopbackRequest(memDescs, fileDescs, false); + } + + TLLM_CUDA_CHECK(cudaMemcpy(memory.data(), cuda_mem, 100, cudaMemcpyDeviceToHost)); + + TLLM_CHECK(memory == fileData); + TLLM_CUDA_CHECK(cudaFree(cuda_mem)); +} + +TEST_P(LoopbackAgentTest, GpuToFile) +{ + std::string const agentName{"loopbackAgent"}; + BaseAgentConfig config{agentName, true, GetParam()}; + auto loopbackAgent = makeLoopbackAgent(config); + + TLLM_CHECK(loopbackAgent); + + std::vector memory(100, 1); + char* cuda_mem; + TLLM_CUDA_CHECK(cudaMalloc(&cuda_mem, 100)); + TLLM_CUDA_CHECK(cudaMemcpy(cuda_mem, memory.data(), 100, cudaMemcpyHostToDevice)); + std::string filename = getDirectory() + std::string("/gpu2file.bin"); + + { + MemoryDesc mem_desc(cuda_mem, 100, 0); + MemoryDescs memDescs{MemoryType::kVRAM, {mem_desc}}; + + std::vector fileDescVec; + fileDescVec.emplace_back(filename, O_CREAT | O_WRONLY, 0664, 100); + FileDescs fileDescs{std::move(fileDescVec)}; + + loopbackAgent->executeLoopbackRequest(memDescs, fileDescs, true); + } + + int fd = ::open(filename.c_str(), O_RDONLY, 0664); + TLLM_CHECK_WITH_INFO(fd >= 0, "Failed to open '%s' for reading", filename.c_str()); + + std::vector fileData(100); + ssize_t bytesRead = ::read(fd, fileData.data(), fileData.size()); + TLLM_CHECK_WITH_INFO(bytesRead == static_cast(fileData.size()), "Failed to read from file"); + ::close(fd); + + TLLM_CHECK(fileData == memory); + TLLM_CUDA_CHECK(cudaFree(cuda_mem)); +} + +INSTANTIATE_TEST_SUITE_P(, LoopbackAgentTest, ::testing::Values(true, false)); From ed5fd420ebe473672f07531c8216c0c826382be7 Mon Sep 17 00:00:00 2001 From: Tomer Shmilovich Date: Wed, 27 Aug 2025 01:36:22 -0700 Subject: [PATCH 7/7] Fix nixl installation on aarch64 CI Signed-off-by: Tomer Shmilovich --- jenkins/Build.groovy | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/jenkins/Build.groovy b/jenkins/Build.groovy index e5f2339f1665..bb480ef8f6df 100644 --- a/jenkins/Build.groovy +++ b/jenkins/Build.groovy @@ -77,12 +77,12 @@ def BUILD_CONFIGS = [ (WHEEL_ARCHS): "80-real;86-real;89-real;90-real;100-real;120-real", ], (CONFIG_LINUX_AARCH64): [ - (WHEEL_EXTRA_ARGS) : "--extra-cmake-vars WARNING_IS_ERROR=ON", + (WHEEL_EXTRA_ARGS) : "--extra-cmake-vars WARNING_IS_ERROR=ON --extra-cmake-vars NIXL_ROOT=/opt/nvidia/nvda_nixl", (TARNAME) : "TensorRT-LLM-GH200.tar.gz", (WHEEL_ARCHS): "90-real;100-real;120-real", ], (CONFIG_LINUX_AARCH64_PYBIND): [ - (WHEEL_EXTRA_ARGS) : "--binding_type pybind --extra-cmake-vars WARNING_IS_ERROR=ON", + (WHEEL_EXTRA_ARGS) : "--binding_type pybind --extra-cmake-vars WARNING_IS_ERROR=ON --extra-cmake-vars NIXL_ROOT=/opt/nvidia/nvda_nixl", (TARNAME) : "pybind-TensorRT-LLM-GH200.tar.gz", (WHEEL_ARCHS): "90-real;100-real;120-real", ],