diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index c0445afc4745..ab7e820e29c1 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -51,6 +51,11 @@ namespace kvc = tensorrt_llm::executor::kv_cache; +namespace tensorrt_llm::batch_manager::kv_cache_manager +{ +class FabricMemory; +} // namespace tensorrt_llm::batch_manager::kv_cache_manager + namespace tensorrt_llm::batch_manager::eviction_policy { class BaseEvictionPolicy; @@ -1344,6 +1349,8 @@ class WindowBlockManager // Buffer manager runtime::BufferManager mBufferManager; + // Fabric memory backing for primary pools (MNNVL-capable allocation) + std::vector> mFabricMemoryPools; // Used to keep track of number of free blocks during scheduling SizeType32 mSchedulingNumFreeBlocks; diff --git a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp index 3be8ee22bc1f..fe47f6c531ad 100644 --- a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp +++ b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp @@ -61,7 +61,7 @@ BaseTransBufferManager::BaseTransBufferManager( mRecvBufferCount = common::getEnvRequestKVCacheConcurrent() ? common::getEnvKVCacheRecvBufferCount() : 1; mSendBufferCount = common::getEnvKVCacheSendMaxConcurrenceNum(); mUseFabricMemory = !(common::getEnvKVCacheTransferUseSyncBuffer() || common::getEnvKVCacheTransferUseAsyncBuffer()) - && kv_cache_manager::FabricMemory::supportFbaricMemory(); + && kv_cache_manager::FabricMemory::supportFabricMemory(); if (mUseFabricMemory) { mTransferBufferSize = kv_cache_manager::FabricMemory::getAlignedSize(mTransferBufferSize); diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp index b6f1e4e8df15..772c9555f0f3 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp @@ -128,7 +128,7 @@ size_t FabricMemory::getAlignedSize(size_t size) return (size + granularity - 1) / granularity * granularity; } -bool FabricMemory::supportFbaricMemory() +bool FabricMemory::supportFabricMemory() { #ifdef __aarch64__ auto support_fun = []() @@ -309,7 +309,7 @@ size_t CacheTransBufferManager::preAllocBufferSize( transferBufferSize += validTokenNum * cacheSizeBytesPerToken; } } - bool useFabricMemory = FabricMemory::supportFbaricMemory() + bool useFabricMemory = FabricMemory::supportFabricMemory() && (!(common::getEnvKVCacheTransferUseSyncBuffer() || common::getEnvKVCacheTransferUseAsyncBuffer())); if (useFabricMemory) { diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.h b/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.h index b63f18ab797f..1635c11bc673 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.h +++ b/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.h @@ -50,7 +50,7 @@ class FabricMemory size_t getSize() const; static size_t getAlignedSize(size_t size); - static bool supportFbaricMemory(); + static bool supportFabricMemory(); private: class Impl; diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index e6c4ff885d1d..c4d265289a71 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -17,12 +17,14 @@ #include "tensorrt_llm/batch_manager/kvCacheManager.h" +#include "tensorrt_llm/batch_manager/cacheTransBuffer.h" #include "tensorrt_llm/batch_manager/common.h" #include "tensorrt_llm/batch_manager/evictionPolicy.h" #include "tensorrt_llm/batch_manager/kvCacheTransferManager.h" #include "tensorrt_llm/batch_manager/radixBlockTree.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/common/memoryUtils.h" #include "tensorrt_llm/executor/executor.h" @@ -1092,6 +1094,23 @@ void WindowBlockManager::allocatePools(bool useUvm) { constexpr nvinfer1::DataType kScaleDtypeNVFP4 = nvinfer1::DataType::kFP8; + bool const requestFabricMemory = tc::getEnvKVCachePoolUseFabricMemory(); + bool const fabricMemorySupported = FabricMemory::supportFabricMemory(); + if (requestFabricMemory && !fabricMemorySupported) + { + TLLM_LOG_WARNING( + "[%s] TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORY=1 was set but fabric memory is not supported on this " + "platform (FabricMemory::supportFabricMemory() returned false); falling back to standard GPU " + "allocation.", + mLogPrefix.c_str()); + } + bool const useFabricMemory = requestFabricMemory && fabricMemorySupported; + + if (useFabricMemory) + { + TLLM_LOG_INFO("[%s] KV cache pool using fabric memory for MNNVL support", mLogPrefix.c_str()); + } + // Allocate a memory pool backing the blocks for each numKvHeads // TODO(oargov): allocate pools in a single buffer and split it, to avoid fragmentation for (auto& pool : mPools) @@ -1124,9 +1143,27 @@ void WindowBlockManager::allocatePools(bool useUvm) cacheShape.d[2], cacheShape.d[3], pool.layerFirstLayout ? " (layer-first)" : ""); if (useUvm) + { pool.primaryPtr = BufferManager::managed(cacheShape, poolDtype); + } + else if (useFabricMemory) + { + auto const numElements = ITensor::volume(cacheShape); + auto const elementSize = tc::getDTypeSize(poolDtype); + auto const totalBytes = static_cast(numElements) * elementSize; + + // Record ownership before exposing the raw pointer: if FabricMemory's ctor throws nothing + // is wrapped; if ITensor::wrap throws afterwards, the unique_ptr in mFabricMemoryPools + // still owns and will free the allocation. + mFabricMemoryPools.reserve(mFabricMemoryPools.size() + 1); + mFabricMemoryPools.emplace_back(std::make_unique(totalBytes)); + pool.primaryPtr = ITensor::wrap(mFabricMemoryPools.back()->getPtr(), poolDtype, cacheShape, numElements); + } else + { pool.primaryPtr = mBufferManager.gpuSync(cacheShape, poolDtype); + } + if (mNumSecondaryBlocks > 0) { nvinfer1::Dims cacheShapeOffload = isRecurrentState() @@ -1149,6 +1186,12 @@ void BlockManager::releasePools() void WindowBlockManager::releasePools() { + if (mTransferManager) + { + mTransferManager->syncTransfers(); + } + mBufferManager.getStream().synchronize(); + for (auto& pool : mPools) { if (pool.primaryPtr) @@ -1160,7 +1203,8 @@ void WindowBlockManager::releasePools() pool.secondaryPtr->release(); } } - mBufferManager.getStream().synchronize(); + // Release fabric memory backing (must happen after ITensor release). + mFabricMemoryPools.clear(); mBufferManager.memoryPoolTrimTo(0); } diff --git a/cpp/tensorrt_llm/batch_manager/rnnCacheTransBuffer.cpp b/cpp/tensorrt_llm/batch_manager/rnnCacheTransBuffer.cpp index 1e52a04b9580..8fa9508cbefe 100644 --- a/cpp/tensorrt_llm/batch_manager/rnnCacheTransBuffer.cpp +++ b/cpp/tensorrt_llm/batch_manager/rnnCacheTransBuffer.cpp @@ -163,7 +163,7 @@ size_t RnnCacheTransBufferManager::preAllocBufferSize( size_t transferBufferSize = rnnStateSizeBytes > 0 ? rnnStateSizeBytes : common::getEnvMemSizeForKVCacheTransferBuffer(); - bool useFabricMemory = kv_cache_manager::FabricMemory::supportFbaricMemory() + bool useFabricMemory = kv_cache_manager::FabricMemory::supportFabricMemory() && (!(common::getEnvKVCacheTransferUseSyncBuffer() || common::getEnvKVCacheTransferUseAsyncBuffer())); if (useFabricMemory) diff --git a/cpp/tensorrt_llm/common/envUtils.cpp b/cpp/tensorrt_llm/common/envUtils.cpp index 61e81c9135af..8987e5048c90 100644 --- a/cpp/tensorrt_llm/common/envUtils.cpp +++ b/cpp/tensorrt_llm/common/envUtils.cpp @@ -493,6 +493,12 @@ bool getEnvKVCacheTransferAllBlocksForWindow() return allBlocksForWindow; } +bool getEnvKVCachePoolUseFabricMemory() +{ + static bool const useFabricMemory = getBoolEnv("TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORY"); + return useFabricMemory; +} + uint16_t getEnvNixlPort() { static uint16_t const nixlPort = getUInt64Env("TRTLLM_NIXL_PORT").value_or(0); diff --git a/cpp/tensorrt_llm/common/envUtils.h b/cpp/tensorrt_llm/common/envUtils.h index 196c15aadbf8..d8e72b7b02cc 100644 --- a/cpp/tensorrt_llm/common/envUtils.h +++ b/cpp/tensorrt_llm/common/envUtils.h @@ -139,6 +139,8 @@ size_t getEnvKVCacheSendMaxConcurrenceNum(); size_t getEnvMemSizeForKVCacheTransferBuffer(); +bool getEnvKVCachePoolUseFabricMemory(); + uint16_t getEnvNixlPort(); bool getEnvNixlEnableCoalesce(); diff --git a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt index eeb5ec34deef..f4e5e9fb2be0 100644 --- a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt +++ b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt @@ -22,6 +22,7 @@ add_gtest(capacitySchedulerTest capacitySchedulerTest.cpp) add_gtest(contextProgressTest contextProgressTest.cu) add_gtest(evictionPolicyTest evictionPolicyTest.cpp) add_gtest(kvCacheManagerTest kvCacheManagerTest.cpp) +add_gtest(kvCacheManagerFabricMemoryTest kvCacheManagerFabricMemoryTest.cpp) add_gtest(kvCacheUtilsTest kvCacheUtilsTest.cpp) add_gtest(llmRequestTest llmRequestTest.cpp) add_gtest(microBatchSchedulerTest microBatchSchedulerTest.cpp) diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerFabricMemoryTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerFabricMemoryTest.cpp new file mode 100644 index 000000000000..d23cbabe4144 --- /dev/null +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerFabricMemoryTest.cpp @@ -0,0 +1,320 @@ +/* + * 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/batch_manager/cacheTransBuffer.h" +#include "tensorrt_llm/batch_manager/common.h" +#include "tensorrt_llm/batch_manager/kvCacheManager.h" +#include "tensorrt_llm/batch_manager/llmRequest.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/runtime/iTensor.h" +#include "tensorrt_llm/runtime/samplingConfig.h" +#include "tensorrt_llm/testing/kvCacheManagerTestUtil.h" + +#include + +#include +#include + +using namespace tensorrt_llm::batch_manager::kv_cache_manager; +using tensorrt_llm::batch_manager::LlmRequest; +namespace tc = tensorrt_llm::common; +namespace tr = tensorrt_llm::runtime; + +using BlocksPerWindow = std::map>; + +// ============================================================================ +// Fabric memory pool allocation tests +// ============================================================================ +// This test binary is separate from kvCacheManagerTest so that setting +// TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORY before the static cache is +// initialized takes effect for all tests. + +namespace +{ + +// Helper: check whether a device pointer was allocated via cuMemCreate with fabric handle. +bool isFabricVmmAllocation(void* ptr) +{ + CUmemGenericAllocationHandle handle{}; + auto ret = cuMemRetainAllocationHandle(&handle, ptr); + if (ret != CUDA_SUCCESS) + { + return false; // Not a VMM allocation (e.g. cudaMalloc) + } + + CUmemAllocationProp prop{}; + auto ret2 = cuMemGetAllocationPropertiesFromHandle(&prop, handle); + cuMemRelease(handle); + if (ret2 != CUDA_SUCCESS) + { + return false; + } + + return (prop.requestedHandleTypes & CU_MEM_HANDLE_TYPE_FABRIC) != 0; +} + +} // namespace + +class KVCacheManagerFabricMemoryTest : public ::testing::Test +{ +protected: + static void SetUpTestSuite() + { + setenv("TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORY", "1", 1); + } + + void SetUp() override + { + if (tc::getDeviceCount() == 0) + { + GTEST_SKIP() << "No CUDA device available"; + } + } +}; + +TEST_F(KVCacheManagerFabricMemoryTest, AllocatePoolsFallbackWhenFabricUnsupported) +{ + if (FabricMemory::supportFabricMemory()) + { + GTEST_SKIP() << "This test targets hardware without fabric memory support"; + } + + auto constexpr numLayers = 4; + auto constexpr numKvHeads = 2; + auto constexpr sizePerHead = 64; + auto constexpr tokensPerBlock = 8; + auto constexpr blocksInPrimaryPool = 16; + auto constexpr blocksInSecondaryPool = 0; + auto constexpr maxNumSequences = 4; + auto constexpr beamWidth = 1; + auto constexpr maxAttentionWindow = tokensPerBlock * blocksInPrimaryPool; + auto const stream = std::make_shared(); + + auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}}; + + BlockManager blockManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, + maxNumSequences, stream, maxAttentionWindow, beamWidth, + std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, 0); + blockManager.allocatePools(false); + + EXPECT_EQ(blockManager.getTokensPerBlock(), tokensPerBlock); + EXPECT_EQ(blockManager.getMaxNumBlocks(), blocksInPrimaryPool); + EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); + + // Pool should use regular GPU allocation, not fabric VMM + auto primaryPool = blockManager.getPrimaryPool(0); + EXPECT_NE(primaryPool, nullptr); + EXPECT_GT(primaryPool->getSize(), 0); + EXPECT_FALSE(isFabricVmmAllocation(primaryPool->data())) + << "Pool should use regular GPU allocation when fabric memory is unsupported"; + + blockManager.releasePools(); +} + +TEST_F(KVCacheManagerFabricMemoryTest, AllocatePoolsWithFabricMemory) +{ + if (!FabricMemory::supportFabricMemory()) + { + GTEST_SKIP() << "Fabric memory not supported on this hardware"; + } + + auto constexpr numLayers = 4; + auto constexpr numKvHeads = 2; + auto constexpr sizePerHead = 64; + auto constexpr tokensPerBlock = 8; + auto constexpr blocksInPrimaryPool = 16; + auto constexpr blocksInSecondaryPool = 0; + auto constexpr maxNumSequences = 4; + auto constexpr beamWidth = 1; + auto constexpr maxAttentionWindow = tokensPerBlock * blocksInPrimaryPool; + auto const stream = std::make_shared(); + + auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}}; + + BlockManager blockManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, + maxNumSequences, stream, maxAttentionWindow, beamWidth, + std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, 0); + blockManager.allocatePools(false); + + EXPECT_EQ(blockManager.getMaxNumBlocks(), blocksInPrimaryPool); + EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); + + // Pool should be a VMM allocation with fabric handle + auto primaryPool = blockManager.getPrimaryPool(0); + EXPECT_NE(primaryPool, nullptr); + EXPECT_GT(primaryPool->getSize(), 0); + EXPECT_TRUE(isFabricVmmAllocation(primaryPool->data())) + << "Pool should be backed by fabric VMM when env var is set and hardware supports it"; + + blockManager.releasePools(); +} + +// Round-trip a fabric-backed primary block through the host-pinned secondary pool to validate +// that DRAM offload/onboard work correctly when the primary is fabric VMM and to make sure +// releasePools() syncs the transfer manager before unmapping the fabric allocation. +TEST_F(KVCacheManagerFabricMemoryTest, OffloadOnboardRoundTripWithFabricPrimary) +{ + if (!FabricMemory::supportFabricMemory()) + { + GTEST_SKIP() << "Fabric memory not supported on this hardware"; + } + + using Element = std::uint16_t; // matches kHALF size + auto constexpr numLayers = 2; + auto constexpr numKvHeads = 2; + auto constexpr sizePerHead = 32; + auto constexpr tokensPerBlock = 4; + auto constexpr blocksInPrimaryPool = 4; + auto constexpr blocksInSecondaryPool = 4; + auto constexpr maxNumSequences = 2; + auto constexpr beamWidth = 1; + auto constexpr beamIdx = 0; + auto constexpr maxAttentionWindow = tokensPerBlock * blocksInPrimaryPool; + auto constexpr transferMode = tensorrt_llm::executor::KvCacheTransferMode::DRAM; + + auto const stream = std::make_shared(); + auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}}; + + BlockManager blockManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, + blocksPerWindow, maxNumSequences, stream, maxAttentionWindow, beamWidth, + std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, 0); + blockManager.allocatePools(false); + + auto primaryPoolPtr = blockManager.getPrimaryPool(0); + auto secondaryPoolPtr = blockManager.getSecondaryPool(0); + ASSERT_NE(primaryPoolPtr, nullptr); + ASSERT_NE(secondaryPoolPtr, nullptr); + EXPECT_TRUE(isFabricVmmAllocation(primaryPoolPtr->data())) + << "Primary pool must be fabric VMM-backed for this test"; + + // Block has shape [2, numLayers, numKvHeads, tokensPerBlock, sizePerHead]. + auto const elemsPerBlock = static_cast(2) * numLayers * numKvHeads * tokensPerBlock * sizePerHead; + auto const bytesPerBlock = elemsPerBlock * sizeof(Element); + + // Drive the primary block allocator through addSequenceBatch. + SizeType32 constexpr maxNewTokens{0}; + tr::SamplingConfig const samplingConfig{beamWidth}; + bool constexpr isStreaming{false}; + + auto inputTokens = std::make_shared(VecTokens{0, 1, 2, 3, 4, 5, 6, 7}); + auto const inputLength = static_cast(inputTokens->size()); + LlmRequest::RequestIdType const requestId{0}; + auto llmRequest = std::make_shared(requestId, maxNewTokens, inputTokens, samplingConfig, isStreaming); + GenerationRequest seq{requestId, inputLength, beamWidth, blockManager.getWindowSizesMetadata()}; + auto const promptLen = llmRequest->getNumTokens(beamIdx); + auto const numContextBlocks = tc::ceilDiv(promptLen, blockManager.getTokensPerBlock()); + auto const seqStats = blockManager.addSequenceBatch({&seq}, {promptLen}, {numContextBlocks}, + {std::ref(*llmRequest)}, maxAttentionWindow, + /*isEnableBlockReuse=*/false); + ASSERT_FALSE(seqStats.empty()); + + auto cacheBlockIds = seq.getCacheBlockIds(maxAttentionWindow).at(beamIdx); + ASSERT_FALSE(cacheBlockIds.empty()); + auto const blockId = cacheBlockIds.front(); + auto block = blockManager.getBlockById(blockId, maxAttentionWindow); + ASSERT_TRUE(block->isPrimary()); + + // Write a known pattern into the fabric-backed primary block. + std::vector hostBuffer(elemsPerBlock); + for (size_t i = 0; i < elemsPerBlock; ++i) + { + hostBuffer[i] = static_cast(i & 0xFFFF); + } + auto const memoryPoolIndex = block->getMemoryPoolBlockIndex(); + auto primaryBlockPtr = tr::ITensor::slice(primaryPoolPtr, memoryPoolIndex, 1); + ASSERT_EQ( + cudaMemcpy(primaryBlockPtr->data(), hostBuffer.data(), bytesPerBlock, cudaMemcpyHostToDevice), cudaSuccess); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + + // Offload fabric primary -> host pinned secondary and verify pattern survived. + blockManager.offloadBlock(block, maxAttentionWindow, transferMode, /*directory=*/""); + EXPECT_FALSE(block->isPrimary()); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + + auto const secondaryIndex = block->getMemoryPoolBlockIndex(); + auto secondaryBlockPtr = tr::ITensor::slice(secondaryPoolPtr, secondaryIndex, 1); + auto const* secondaryRaw = reinterpret_cast(secondaryBlockPtr->data()); + for (size_t i = 0; i < elemsPerBlock; ++i) + { + ASSERT_EQ(secondaryRaw[i], hostBuffer[i]) << "Mismatch at element " << i << " after fabric->host offload"; + } + + // Onboard host secondary -> fabric primary, then read back and verify integrity. + blockManager.onboardBlock(seq, block, maxAttentionWindow, transferMode, /*directory=*/""); + EXPECT_TRUE(block->isPrimary()); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + + std::vector readback(elemsPerBlock, 0); + auto const onboardIndex = block->getMemoryPoolBlockIndex(); + auto reboardBlockPtr = tr::ITensor::slice(primaryPoolPtr, onboardIndex, 1); + ASSERT_EQ(cudaMemcpy(readback.data(), reboardBlockPtr->data(), bytesPerBlock, cudaMemcpyDeviceToHost), cudaSuccess); + for (size_t i = 0; i < elemsPerBlock; ++i) + { + ASSERT_EQ(readback[i], hostBuffer[i]) << "Mismatch at element " << i << " after host->fabric onboard"; + } + + tensorrt_llm::testing::KvCacheManagerTestUtil::simulatePrefillCompletion(*llmRequest); + blockManager.releaseBlocks(seq, llmRequest); + + // releasePools() must syncTransfers() + sync the buffer-manager stream before unmapping the + // fabric allocation; if it does not, the transfer queued above can race with cuMemUnmap. + blockManager.releasePools(); +} + +TEST_F(KVCacheManagerFabricMemoryTest, ReleasePoolsClearsFabricMemory) +{ + auto constexpr numLayers = 2; + auto constexpr numKvHeads = 2; + auto constexpr sizePerHead = 64; + auto constexpr tokensPerBlock = 4; + auto constexpr blocksInPrimaryPool = 8; + auto constexpr blocksInSecondaryPool = 0; + auto constexpr maxNumSequences = 2; + auto constexpr beamWidth = 1; + auto constexpr maxAttentionWindow = tokensPerBlock * blocksInPrimaryPool; + auto const stream = std::make_shared(); + + auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}}; + + BlockManager blockManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, + maxNumSequences, stream, maxAttentionWindow, beamWidth, + std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, 0); + + size_t freeBefore = 0; + size_t freeAfterAlloc = 0; + size_t freeAfterRelease = 0; + size_t total = 0; + + cudaMemGetInfo(&freeBefore, &total); + + blockManager.allocatePools(false); + cudaMemGetInfo(&freeAfterAlloc, &total); + EXPECT_LT(freeAfterAlloc, freeBefore) << "allocatePools should consume GPU memory"; + + blockManager.releasePools(); + cudaMemGetInfo(&freeAfterRelease, &total); + EXPECT_GE(freeAfterRelease, freeBefore - (1U << 20U)) + << "releasePools should return GPU memory (within 1MB tolerance)"; + + // Second cycle: verify no accumulation + cudaMemGetInfo(&freeBefore, &total); + blockManager.allocatePools(false); + EXPECT_EQ(blockManager.getMaxNumBlocks(), blocksInPrimaryPool); + blockManager.releasePools(); + cudaMemGetInfo(&freeAfterRelease, &total); + EXPECT_GE(freeAfterRelease, freeBefore - (1U << 20U)) << "Second release should also return GPU memory"; +} diff --git a/tests/integration/defs/disaggregated/test_disaggregated.py b/tests/integration/defs/disaggregated/test_disaggregated.py index 7223f1b47fbb..5f2cfb0f2c86 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated.py +++ b/tests/integration/defs/disaggregated/test_disaggregated.py @@ -1130,6 +1130,30 @@ def test_disaggregated_overlap_transceiver_runtime_python( cwd=llm_venv.get_working_directory()) +# Exercises the disaggregated KV-cache transfer path with the Python cache transceiver runtime +# while the KV-cache pool itself is allocated from fabric (MNNVL) VMM memory via +# TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORY=1. Restricted to GB200/GB300 since those are the only +# platforms with MNNVL fabric-memory support; on other devices the env var would silently fall +# back to a non-fabric allocation, which would defeat the purpose of this test. +@pytest.mark.skip_device_not_contain(["GB200", "GB300"]) +@pytest.mark.parametrize("llama_model_root", ['TinyLlama-1.1B-Chat-v1.0'], + indirect=True) +def test_disaggregated_overlap_transceiver_runtime_python_fabric_memory( + disaggregated_test_root, llm_venv, disaggregated_example_root, + llama_model_root): + setup_model_symlink(llm_venv, llama_model_root, + "TinyLlama/TinyLlama-1.1B-Chat-v1.0") + + env = llm_venv._new_env.copy() + env["UCX_TLS"] = get_ucx_tls() + env["TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORY"] = "1" + run_disaggregated_test(disaggregated_example_root, + "overlap_transceiver_runtime_python", + env=env, + model_path=llama_model_root, + cwd=llm_venv.get_working_directory()) + + @pytest.mark.parametrize("llama_model_root", ['TinyLlama-1.1B-Chat-v1.0'], indirect=True) def test_disaggregated_perf_metrics(disaggregated_test_root, llm_venv, diff --git a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml index e90d96ddef20..f12e16220744 100644 --- a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml +++ b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml @@ -65,6 +65,7 @@ l0_gb200_multi_gpus: - accuracy/test_dwdp_disaggregated_serving.py::TestDwdpDeepSeekV3Lite::test_dwdp_accuracy_mode_b_overlap - unittest/_torch/modules/moe/test_moe_comm.py::TestMoEComm::test_moe_comm - unittest/_torch/modules/moe/test_moe_comm.py::TestMoEComm::test_moe_comm_postquant + - disaggregated/test_disaggregated.py::test_disaggregated_overlap_transceiver_runtime_python_fabric_memory[TinyLlama-1.1B-Chat-v1.0] - condition: ranges: