mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-01-14 06:27:45 +08:00
285 lines
8.9 KiB
C++
285 lines
8.9 KiB
C++
/*
|
|
* Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved.
|
|
*
|
|
* 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 "bufferManager.h"
|
|
#include "cudaMemPool.h"
|
|
#include "tensorrt_llm/common/assert.h"
|
|
#include "tensorrt_llm/common/cudaUtils.h"
|
|
#include "tensorrt_llm/common/memoryUtils.h"
|
|
#include "tllmBuffers.h"
|
|
|
|
#include <cstring>
|
|
#include <cuda_runtime_api.h>
|
|
#include <memory>
|
|
|
|
namespace tc = tensorrt_llm::common;
|
|
|
|
namespace tensorrt_llm::runtime
|
|
{
|
|
|
|
BufferManager::BufferManager(CudaStreamPtr stream, bool trimPool)
|
|
: mStream{std::move(stream)}
|
|
, mTrimPool{trimPool}
|
|
{
|
|
TLLM_CHECK_WITH_INFO(static_cast<bool>(mStream), "Undefined CUDA stream");
|
|
mPool = CudaMemPool::getPrimaryPoolForDevice(mStream->getDevice());
|
|
}
|
|
|
|
BufferManager::IBufferPtr BufferManager::gpu(std::size_t size, nvinfer1::DataType type) const
|
|
{
|
|
if (auto vmAllocator = getVirtualMemoryAllocator())
|
|
{
|
|
return std::make_unique<VirtualAddressDeviceBuffer>(size, type, std::move(vmAllocator));
|
|
}
|
|
if (static_cast<bool>(mPool))
|
|
{
|
|
return std::make_unique<DeviceBuffer>(size, type, CudaAllocatorAsync{mStream, mPool});
|
|
}
|
|
// When memory pools are not supported, fallback to synchronous memory allocations.
|
|
return gpuSync(size, type);
|
|
}
|
|
|
|
BufferManager::ITensorPtr BufferManager::gpu(nvinfer1::Dims dims, nvinfer1::DataType type) const
|
|
{
|
|
if (auto vmAllocator = getVirtualMemoryAllocator())
|
|
{
|
|
return std::make_unique<VirtualAddressDeviceTensor>(dims, type, std::move(vmAllocator));
|
|
}
|
|
if (static_cast<bool>(mPool))
|
|
{
|
|
return std::make_unique<DeviceTensor>(dims, type, CudaAllocatorAsync{mStream, mPool});
|
|
}
|
|
// When memory pools are not supported, fallback to synchronous memory allocations.
|
|
return gpuSync(dims, type);
|
|
}
|
|
|
|
BufferManager::IBufferPtr BufferManager::gpuSync(std::size_t size, nvinfer1::DataType type)
|
|
{
|
|
if (auto vmAllocator = getVirtualMemoryAllocator())
|
|
{
|
|
return std::make_unique<VirtualAddressDeviceBuffer>(size, type, std::move(vmAllocator));
|
|
}
|
|
return std::make_unique<StaticDeviceBuffer>(size, type, CudaAllocator{});
|
|
}
|
|
|
|
BufferManager::ITensorPtr BufferManager::gpuSync(nvinfer1::Dims dims, nvinfer1::DataType type)
|
|
{
|
|
if (auto vmAllocator = getVirtualMemoryAllocator())
|
|
{
|
|
return std::make_unique<VirtualAddressDeviceTensor>(dims, type, std::move(vmAllocator));
|
|
}
|
|
return std::make_unique<StaticDeviceTensor>(dims, type, CudaAllocator{});
|
|
}
|
|
|
|
BufferManager::IBufferPtr BufferManager::cpu(std::size_t size, nvinfer1::DataType type)
|
|
{
|
|
return std::make_unique<HostBuffer>(size, type);
|
|
}
|
|
|
|
BufferManager::ITensorPtr BufferManager::cpu(nvinfer1::Dims dims, nvinfer1::DataType type)
|
|
{
|
|
return std::make_unique<HostTensor>(dims, type);
|
|
}
|
|
|
|
BufferManager::IBufferPtr BufferManager::pinned(std::size_t size, nvinfer1::DataType type)
|
|
{
|
|
return std::make_unique<PinnedBuffer>(size, type);
|
|
}
|
|
|
|
BufferManager::ITensorPtr BufferManager::pinned(nvinfer1::Dims dims, nvinfer1::DataType type)
|
|
{
|
|
return std::make_unique<PinnedTensor>(dims, type);
|
|
}
|
|
|
|
BufferManager::IBufferPtr BufferManager::pinnedPool(std::size_t size, nvinfer1::DataType type)
|
|
{
|
|
return std::make_unique<PinnedPoolBuffer>(size, type);
|
|
}
|
|
|
|
BufferManager::ITensorPtr BufferManager::pinnedPool(nvinfer1::Dims dims, nvinfer1::DataType type)
|
|
{
|
|
return std::make_unique<PinnedPoolTensor>(dims, type);
|
|
}
|
|
|
|
BufferManager::IBufferPtr BufferManager::managed(std::size_t size, nvinfer1::DataType type)
|
|
{
|
|
return std::make_unique<UVMBuffer>(size, type);
|
|
}
|
|
|
|
BufferManager::ITensorPtr BufferManager::managed(nvinfer1::Dims dims, nvinfer1::DataType type)
|
|
{
|
|
return std::make_unique<UVMTensor>(dims, type);
|
|
}
|
|
|
|
BufferManager::ITensorPtr BufferManager::ipcNvls(std::set<int> ranks, nvinfer1::Dims dims, nvinfer1::DataType type)
|
|
{
|
|
return std::make_unique<MulticastTensor>(dims, type, ranks);
|
|
}
|
|
|
|
void BufferManager::setZero(IBuffer& buffer) const
|
|
{
|
|
setMem(buffer, 0);
|
|
}
|
|
|
|
void BufferManager::setMem(IBuffer& buffer, int32_t value) const
|
|
{
|
|
if (buffer.getMemoryType() == MemoryType::kGPU)
|
|
{
|
|
TLLM_CUDA_CHECK(cudaMemsetAsync(buffer.data(), value, buffer.getSizeInBytes(), mStream->get()));
|
|
}
|
|
else
|
|
{
|
|
std::memset(buffer.data(), value, buffer.getSizeInBytes());
|
|
}
|
|
}
|
|
|
|
void BufferManager::copy(void const* src, IBuffer& dst, MemoryType srcType) const
|
|
{
|
|
if (dst.getSizeInBytes() > 0)
|
|
{
|
|
if (srcType != MemoryType::kGPU && dst.getMemoryType() != MemoryType::kGPU)
|
|
{
|
|
std::memcpy(dst.data(), src, dst.getSizeInBytes());
|
|
}
|
|
else
|
|
{
|
|
TLLM_CUDA_CHECK(tensorrt_llm::common::cudaMemcpyAsyncSanitized(
|
|
dst.data(), src, dst.getSizeInBytes(), cudaMemcpyDefault, mStream->get()));
|
|
}
|
|
}
|
|
}
|
|
|
|
void BufferManager::copy(IBuffer const& src, void* dst, MemoryType dstType) const
|
|
{
|
|
if (src.getSizeInBytes() > 0)
|
|
{
|
|
if (src.getMemoryType() != MemoryType::kGPU && dstType != MemoryType::kGPU)
|
|
{
|
|
std::memcpy(dst, src.data(), src.getSizeInBytes());
|
|
}
|
|
else
|
|
{
|
|
TLLM_CUDA_CHECK(tensorrt_llm::common::cudaMemcpyAsyncSanitized(
|
|
dst, src.data(), src.getSizeInBytes(), cudaMemcpyDefault, mStream->get()));
|
|
}
|
|
}
|
|
}
|
|
|
|
void BufferManager::copy(IBuffer const& src, IBuffer& dst) const
|
|
{
|
|
TLLM_CHECK_WITH_INFO(src.getDataType() == dst.getDataType(),
|
|
tc::fmtstr("Incompatible data types: %s != %s", src.getDataTypeName(), dst.getDataTypeName()));
|
|
TLLM_CHECK_WITH_INFO(src.getSizeInBytes() == dst.getSizeInBytes(),
|
|
tc::fmtstr("Incompatible buffer sizes: %lu != %lu", src.getSizeInBytes(), dst.getSizeInBytes()));
|
|
copy(src, dst.data(), dst.getMemoryType());
|
|
}
|
|
|
|
BufferManager::IBufferPtr BufferManager::allocate(
|
|
MemoryType memoryType, std::size_t size, nvinfer1::DataType type) const
|
|
{
|
|
switch (memoryType)
|
|
{
|
|
case MemoryType::kCPU: return cpu(size, type);
|
|
case MemoryType::kGPU: return gpu(size, type);
|
|
case MemoryType::kPINNED: return pinned(size, type);
|
|
case MemoryType::kUVM: return managed(size, type);
|
|
case MemoryType::kPINNEDPOOL: return pinnedPool(size, type);
|
|
}
|
|
|
|
TLLM_THROW("Unknown memory type");
|
|
}
|
|
|
|
BufferManager::ITensorPtr BufferManager::allocate(
|
|
MemoryType memoryType, nvinfer1::Dims dims, nvinfer1::DataType type) const
|
|
{
|
|
switch (memoryType)
|
|
{
|
|
case MemoryType::kCPU: return cpu(dims, type);
|
|
case MemoryType::kGPU: return gpu(dims, type);
|
|
case MemoryType::kPINNED: return pinned(dims, type);
|
|
case MemoryType::kUVM: return managed(dims, type);
|
|
case MemoryType::kPINNEDPOOL: return pinnedPool(dims, type);
|
|
}
|
|
|
|
TLLM_THROW("Unknown memory type");
|
|
}
|
|
|
|
BufferManager::IBufferPtr BufferManager::copyFrom(IBuffer const& src, MemoryType memoryType) const
|
|
{
|
|
auto dst = allocate(memoryType, src.getSize(), src.getDataType());
|
|
copy(src, *dst);
|
|
return dst;
|
|
}
|
|
|
|
BufferManager::ITensorPtr BufferManager::copyFrom(ITensor const& src, MemoryType memoryType) const
|
|
{
|
|
auto dst = allocate(memoryType, src.getShape(), src.getDataType());
|
|
copy(src, *dst);
|
|
return dst;
|
|
}
|
|
|
|
CudaStream const& BufferManager::getStream() const
|
|
{
|
|
return *mStream;
|
|
}
|
|
|
|
std::size_t BufferManager::memoryPoolReserved() const
|
|
{
|
|
if (!static_cast<bool>(mPool))
|
|
{
|
|
TLLM_LOG_TRACE(
|
|
"Operation '%s' trivially returns zero on systems without memory pool support.", __PRETTY_FUNCTION__);
|
|
return 0;
|
|
}
|
|
mStream->synchronize();
|
|
return mPool->memoryPoolReserved();
|
|
}
|
|
|
|
std::size_t BufferManager::memoryPoolUsed() const
|
|
{
|
|
if (!static_cast<bool>(mPool))
|
|
{
|
|
TLLM_LOG_TRACE(
|
|
"Operation '%s' trivially returns zero on systems without memory pool support.", __PRETTY_FUNCTION__);
|
|
return 0;
|
|
}
|
|
mStream->synchronize();
|
|
return mPool->memoryPoolUsed();
|
|
}
|
|
|
|
std::size_t BufferManager::memoryPoolFree() const
|
|
{
|
|
if (!static_cast<bool>(mPool))
|
|
{
|
|
TLLM_LOG_TRACE(
|
|
"Operation '%s' trivially returns zero on systems without memory pool support.", __PRETTY_FUNCTION__);
|
|
return 0;
|
|
}
|
|
mStream->synchronize();
|
|
return mPool->memoryPoolReserved() - mPool->memoryPoolUsed();
|
|
}
|
|
|
|
void BufferManager::memoryPoolTrimTo(std::size_t size)
|
|
{
|
|
if (!static_cast<bool>(mPool))
|
|
{
|
|
TLLM_LOG_TRACE("Operation '%s' does not do anything on this system as it does not support memory pools.",
|
|
__PRETTY_FUNCTION__);
|
|
return;
|
|
}
|
|
mPool->memoryPoolTrimTo(size);
|
|
}
|
|
} // namespace tensorrt_llm::runtime
|