TllmRuntime

class TllmRuntime{public:    using TensorMap = StringPtrMap<ITensor>;explicit TllmRuntime(void const* engineData, std::size_t engineSize, nvinfer1::ILogger& logger);explicit TllmRuntime(nvinfer1::IHostMemory const& engineBuffer, nvinfer1::ILogger& logger)        : TllmRuntime{engineBuffer.data(), engineBuffer.size(), logger}    {    }explicit TllmRuntime(void const* engineData, std::size_t engineSize);explicit TllmRuntime(nvinfer1::IHostMemory const& engineBuffer)        : TllmRuntime{engineBuffer.data(), engineBuffer.size()}    {    }    SizeType getNbContexts() const{        return static_cast<SizeType>(mContexts.size());    }    nvinfer1::IExecutionContext& getContext(SizeType contextIndex) const{        return *mContexts.at(contextIndex);    }    SizeType getNbProfiles() const{        return static_cast<SizeType>(mEngine->getNbOptimizationProfiles());    }    nvinfer1::IExecutionContext& addContext(std::int32_t profileIndex);void clearContexts();void setInputTensors(SizeType contextIndex, TensorMap const& tensorMap);void setOutputTensors(SizeType contextIndex, TensorMap& tensorMap);bool executeContext(SizeType contextIndex) const;    CudaStream const& getStream() const;    BufferManager::CudaStreamPtr getStreamPtr(){        return mStream;    }    nvinfer1::ICudaEngine& getEngine(){        return *mEngine;    }    nvinfer1::ICudaEngine const& getEngine() const{        return *mEngine;    }    BufferManager& getBufferManager(){        return mBufferManager;    }    BufferManager const& getBufferManager() const{        return mBufferManager;    }private:    BufferManager::CudaStreamPtr mStream;    BufferManager mBufferManager;    std::unique_ptr<nvinfer1::IRuntime> mRuntime;    std::unique_ptr<nvinfer1::ICudaEngine> mEngine;    BufferManager::IBufferPtr mEngineBuffer;    std::vector<std::unique_ptr<nvinfer1::IExecutionContext>> mContexts;    std::unique_ptr<ITensor> mDummyTensor;};

TllmRuntime::TllmRuntime(void const* engineData, std::size_t engineSize, nvinfer1::ILogger& logger)    : mStream(std::make_shared<CudaStream>())    , mBufferManager{mStream}    , mRuntime{nvinfer1::createInferRuntime(logger)}    , mEngine{mRuntime->deserializeCudaEngine(engineData, engineSize)}{    TLLM_CHECK_WITH_INFO(mEngine != nullptr, "Failed to deserialize cuda engine");    auto const devMemorySize = mEngine->getDeviceMemorySize();    mEngineBuffer = mBufferManager.gpu(devMemorySize);}TllmRuntime::TllmRuntime(void const* engineData, std::size_t engineSize)    : TllmRuntime{engineData, engineSize, defaultLogger}{}CudaStream const& TllmRuntime::getStream() const{    return *mStream;}

namespace{using DimType = std::remove_reference_t<decltype(std::declval<nvinfer1::Dims>().d[0])>;static_assert(sizeof(SizeType) >= sizeof(DimType), "SizeType is too small");static_assert(std::is_signed<SizeType>::value, "SizeType must be signed");nvinfer1::Dims shapeToDims(std::vector<std::size_t> const& shape){    TLLM_CHECK(shape.size() <= nvinfer1::Dims::MAX_DIMS);    nvinfer1::Dims dims;    auto constexpr dim_max = std::numeric_limits<DimType>::max();    dims.nbDims = static_cast<std::int32_t>(shape.size());    for (std::size_t i = 0; i < shape.size(); ++i)    {        // shape[i] >= 0 because it has unsigned type. Check upper bound:        TLLM_CHECK(shape[i] <= static_cast<std::size_t>(dim_max));        dims.d[i] = static_cast<DimType>(shape[i]);    }    return dims;}std::vector<std::size_t> dimsToShape(nvinfer1::Dims const& dims){    TLLM_CHECK(dims.nbDims >= 0);    std::vector<std::size_t> shape(dims.nbDims);    for (std::int32_t i = 0; i < dims.nbDims; ++i)    {        TLLM_CHECK(dims.d[i] >= 0);        shape[i] = static_cast<std::size_t>(dims.d[i]);    }    return shape;}tensorrt_llm::runtime::TllmLogger defaultLogger{};} // namespace

addContext

nvinfer1::IExecutionContext& TllmRuntime::addContext(std::int32_t profileIndex){    TLLM_CHECK(0 <= profileIndex && profileIndex < mEngine->getNbOptimizationProfiles());    mContexts.emplace_back(mEngine->createExecutionContextWithoutDeviceMemory());    auto& context = *mContexts.back();    context.setDeviceMemory(mEngineBuffer->data());    context.setOptimizationProfileAsync(profileIndex, mStream->get());    return context;}

clearContexts

void TllmRuntime::clearContexts(){    for (auto& context : mContexts)    {        context.reset();    }    mContexts.clear();}

executeContext

bool TllmRuntime::executeContext(SizeType contextIndex) const{    NVTX3_FUNC_RANGE();    auto& context = getContext(contextIndex);    return context.enqueueV3(mStream->get());}

setInputTensors

void TllmRuntime::setInputTensors(SizeType contextIndex, TensorMap const& tensorMap){    NVTX3_FUNC_RANGE();    TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__);    auto& context = getContext(contextIndex);    for (std::int32_t i = 0; i < mEngine->getNbIOTensors(); ++i)    {        auto const name = mEngine->getIOTensorName(i);        if (mEngine->getTensorIOMode(name) == nvinfer1::TensorIOMode::kINPUT)        {            NVTX3_SCOPED_RANGE(input_tensor);            auto pos = tensorMap.find(name);            if (pos == tensorMap.end())            {                auto expectedShape = mEngine->getTensorShape(name);                TLLM_THROW(                    "Input tensor '%s' not found; expected shape: %s", name, ITensor::toString(expectedShape).c_str());            }            auto const& tensor = pos->second;            auto const tensorDtype = tensor->getDataType();            auto const engineDtype = mEngine->getTensorDataType(name);            // WAR: TRT does not support mixed FP8 and FP16 input, so engine expects FP16 tensors.            TLLM_CHECK_WITH_INFO(tensorDtype == engineDtype                    || (tensorDtype == nvinfer1::DataType::kFP8 && engineDtype == nvinfer1::DataType::kHALF),                tc::fmtstr("%s: expected type %d, provided type %d", name, static_cast<std::int32_t>(engineDtype),                    static_cast<std::int32_t>(tensorDtype)));            auto const shapeExpected = mEngine->getTensorShape(name);            auto const shapeProvided = tensor->getShape();            TLLM_CHECK_WITH_INFO(shapeExpected.nbDims == shapeProvided.nbDims,                tc::fmtstr("%s: expected %d dims, provided %d dims", name, shapeExpected.nbDims, shapeProvided.nbDims));            for (SizeType j = 0; j < shapeExpected.nbDims; ++j)            {                auto const dimExpected = shapeExpected.d[j];                auto const dimProvided = shapeProvided.d[j];                if (dimExpected >= 0 && dimExpected != dimProvided)                {                    TLLM_LOG_WARNING(                        "%s: expected dim[%d] = %d, provided dim[%d] = %d", name, j, dimExpected, j, dimProvided);                }            }            TLLM_CHECK_WITH_INFO(context.setInputShape(name, shapeProvided), name);            auto* const data = tensor->data();            if (data)            {                context.setInputTensorAddress(name, data);            }            else            {                TLLM_CHECK_WITH_INFO(tensor->getSize() == 0, std::string("Invalid data for tensor: ") + name);                // TensorRT runtime does not support nullptr.                if (!mDummyTensor)                {                    mDummyTensor = mBufferManager.gpu(ITensor::makeShape({1}));                }                context.setInputTensorAddress(name, mDummyTensor->data());            }        }    }    {        NVTX3_SCOPED_RANGE(infer_shapes);        char const* missing;        auto const nbMissing = context.inferShapes(1, &missing);        if (nbMissing > 0)        {            TLLM_THROW("Input shape not specified: %s", missing);        }        else if (nbMissing < 0)        {            TLLM_THROW("Invalid input shape");        }    }    {        NVTX3_SCOPED_RANGE(final_checks);        TLLM_CHECK_WITH_INFO(context.allInputDimensionsSpecified(), "Input dimensions not specified");        TLLM_CHECK_WITH_INFO(context.allInputShapesSpecified(), "Input shapes not specified");    }}

setOutputTensors

void TllmRuntime::setOutputTensors(SizeType contextIndex, TensorMap& tensorMap){    NVTX3_FUNC_RANGE();    auto& context = getContext(contextIndex);    for (std::int32_t i = 0; i < mEngine->getNbIOTensors(); ++i)    {        auto const name = mEngine->getIOTensorName(i);        if (mEngine->getTensorIOMode(name) == nvinfer1::TensorIOMode::kOUTPUT)        {            NVTX3_SCOPED_RANGE(output_tensor);            auto const dims = context.getTensorShape(name);            auto const engineDtype = mEngine->getTensorDataType(name);            auto pos = tensorMap.find(name);            if (pos != tensorMap.end())            {                auto const& tensor = pos->second;                auto const tensorDtype = tensor->getDataType();                // WAR: TRT does not support mixed FP8 and FP16 input, so engine expects FP16 tensors.                TLLM_CHECK_WITH_INFO(tensorDtype == engineDtype                        || (tensorDtype == nvinfer1::DataType::kFP8 && engineDtype == nvinfer1::DataType::kHALF),                    tc::fmtstr("%s: expected type %d, provided type %d", name, static_cast<std::int32_t>(engineDtype),                        static_cast<std::int32_t>(tensorDtype)));                tensor->reshape(dims);                context.setTensorAddress(name, tensor->data());            }            else            {                auto tensor = ITensor::SharedPtr(mBufferManager.gpu(dims, engineDtype));                tensorMap.insert(pos, std::make_pair(name, tensor));                context.setTensorAddress(name, tensor->data());            }        }    }}

torch.h

class Torch{public:static at::Tensor tensor(ITensor::SharedPtr tensor){        auto const tensorOptions = at::device(TorchUtils::device((*tensor).data()))                                       .pinned_memory((*tensor).getMemoryType() == MemoryType::kPINNED)                                       .dtype(TorchUtils::dataType((*tensor).getDataType()))                                       .layout(at::kStrided);        return at::for_blob(tensor->data(), TorchUtils::shape(tensor->getShape())) // NOLINT(*-use-after-move)            .options(tensorOptions)            .deleter(                [ptr = std::move(tensor)](void* data) mutable                {                    try                    {                        TLLM_CHECK(data == ptr->data());                        ptr.reset();                    }                    catch (std::exception const& e)                    {                        TLLM_LOG_EXCEPTION(e);                    }                })            .make_tensor();    }static at::Tensor buffer(IBuffer::SharedPtr buffer){        auto const shape = ITensor::makeShape({static_cast<runtime::SizeType>(buffer->getSize())});        return tensor(ITensor::view(std::move(buffer), shape));    }static void setCurrentStream(runtime::CudaStream& cudaStream){        at::cuda::setCurrentCUDAStream(TorchUtils::stream(cudaStream));    }private:    Torch() = default;};

torchUtils.h

class TorchUtils{public:    using SizeType = at::IntArrayRef::value_type;static std::vector<SizeType> shape(ITensor::Shape const& dims){        TLLM_CHECK(dims.nbDims >= 0);        std::vector<SizeType> shape{};        shape.reserve(dims.nbDims);        std::transform(            dims.d, dims.d + dims.nbDims, std::back_inserter(shape), [](auto x) { return static_cast<SizeType>(x); });        return shape;    }static ITensor::Shape shape(at::IntArrayRef const& sizes){        TLLM_CHECK(sizes.size() <= ITensor::Shape::MAX_DIMS);        ITensor::Shape shape{static_cast<runtime::SizeType>(sizes.size())};        using dimType = std::remove_reference_t<decltype(shape.d[0])>;        for (std::size_t i = 0; i < sizes.size(); ++i)        {            TLLM_CHECK(sizes[i] <= std::numeric_limits<dimType>::max());            shape.d[i] = static_cast<dimType>(sizes[i]);        }        return shape;    }static std::vector<SizeType> makeShape(std::initializer_list<runtime::SizeType> sizes){        return shape(ITensor::makeShape(sizes));    }static at::Device device(void const* ptr){        ::cudaPointerAttributes attr{};        TLLM_CUDA_CHECK(cudaPointerGetAttributes(&attr, ptr));        auto const memoryType = attr.type;        return (memoryType == ::cudaMemoryTypeDevice || memoryType == ::cudaMemoryTypeManaged)            ? at::Device{at::kCUDA, static_cast<at::DeviceIndex>(attr.device)}            : at::Device{at::kCPU};    }static at::ScalarType dataType(IBuffer::DataType dataType){        switch (dataType)        {        case IBuffer::DataType::kFLOAT: return at::ScalarType::Float;        case IBuffer::DataType::kHALF: return at::ScalarType::Half;        case IBuffer::DataType::kINT8: return torch::kInt8;        case IBuffer::DataType::kUINT8: return torch::kUInt8;        case IBuffer::DataType::kINT32: return torch::kInt32;        case IBuffer::DataType::kINT64: return torch::kInt64;        case IBuffer::DataType::kBOOL: return at::ScalarType::Bool;        case IBuffer::DataType::kFP8: return at::ScalarType::Bits8;        case IBuffer::DataType::kBF16: return at::ScalarType::BFloat16;        default: TLLM_THROW("unsupported data type");        }    }static IBuffer::DataType dataType(at::ScalarType scalarType){        switch (scalarType)        {        case at::ScalarType::Float: return IBuffer::DataType::kFLOAT;        case at::ScalarType::Half: return IBuffer::DataType::kHALF;        case torch::kInt8: return IBuffer::DataType::kINT8;        case torch::kUInt8: return IBuffer::DataType::kUINT8;        case torch::kInt32: return IBuffer::DataType::kINT32;        case torch::kInt64: return IBuffer::DataType::kINT64;        case at::ScalarType::Bool: return IBuffer::DataType::kBOOL;        case at::ScalarType::Bits8: return IBuffer::DataType::kFP8;        case at::ScalarType::BFloat16: return IBuffer::DataType::kBF16;        default: TLLM_THROW("unsupported data type");        }    }    static at::cuda::CUDAStream stream(runtime::CudaStream& cudaStream){        return at::cuda::getStreamFromExternal(cudaStream.get(), static_cast<at::DeviceIndex>(cudaStream.getDevice()));    }private:    TorchUtils() = default;};

torchView.h

class TorchView : virtual public ITensor{public:static ITensor::UniquePtr of(at::Tensor&& tensor){        return ITensor::UniquePtr{new TorchView{std::move(tensor)}};    }static ITensor::UniquePtr of(at::Tensor tensor){        return ITensor::UniquePtr{new TorchView{std::move(tensor)}};    }void* data() override{        if (getSize() == 0)            return nullptr;        return mTensor.data_ptr();    }    [[nodiscard]]void const* data() const override{        if (getSize() == 0)            return nullptr;        return mTensor.data_ptr();    }    [[nodiscard]]size_t getSize() const override{        return mTensor.numel();    }    [[nodiscard]] std::size_t getCapacity() const override{        return mCapacity;    }    [[nodiscard]] DataType getDataType() const override{        return TorchUtils::dataType(mTensor.scalar_type());    }    [[nodiscard]] MemoryType getMemoryType() const override{        return mTensor.is_cuda() ? MemoryType::kGPU : mTensor.is_pinned() ? MemoryType::kPINNED : MemoryType::kCPU;    }void resize(std::size_t newSize) override{        TLLM_CHECK(newSize <= getCapacity());        if (newSize != getSize())        {            using dimType = std::remove_reference_t<decltype(mDims.d[0])>;            auto constexpr max_size = std::numeric_limits<dimType>::max();            TLLM_CHECK_WITH_INFO(newSize <= max_size, "New size is too large. Use reshape() instead.");            mTensor.resize_({static_cast<at::IntArrayRef::value_type>(newSize)});            mDims.nbDims = 1;            mDims.d[0] = static_cast<dimType>(newSize);        }    }void release() override{        resize(0);    }    [[nodiscard]] Shape const& getShape() const override{        return mDims;    }void reshape(Shape const& dims) override{        TLLM_CHECK(volumeNonNegative(dims) <= getCapacity());        mTensor.resize_(TorchUtils::shape(dims));        mDims = dims;    }private:explicit TorchView(at::Tensor&& tensor)        : mTensor(tensor)        , mDims{TorchUtils::shape(mTensor.sizes())}        , mCapacity{static_cast<std::size_t>(mTensor.numel())}    {        TLLM_CHECK(mTensor.is_contiguous());    };    at::Tensor mTensor;    Shape mDims;    std::size_t mCapacity;};

参考文献

• • https://github.com/NVIDIA/TensorRT-LLM/blob/release/0.5.0/cpp/tensorrt_llm/runtime/tllmRuntime.h