diff --git a/.gitignore b/.gitignore index 7addd122..5885f9e2 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,6 @@ build/ +build-plugin-trt10/ +build-plugin-trt11/ dist/ *.egg-info *_engine/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index ca9ad0fe..7a95968e 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -63,7 +63,7 @@ repos: hooks: - id: codespell args: - - --ignore-words-list=rouge,inout,atleast,strat,nd,subtile,thrid,ans,datas,thw + - --ignore-words-list=rouge,inout,atleast,strat,nd,subtile,thrid,ans,datas,thw,mot exclude: '.*\.json$' - repo: https://github.com/PyCQA/autoflake rev: v2.3.3 diff --git a/README.md b/README.md index 9c46449a..c302230a 100644 --- a/README.md +++ b/README.md @@ -15,7 +15,7 @@ ## Overview -TensorRT Edge-LLM is NVIDIA's high-performance C++ inference runtime for Large Language Models (LLMs) and Vision-Language Models (VLMs) on embedded platforms. It enables efficient deployment of state-of-the-art language models on resource-constrained devices such as NVIDIA Jetson, NVIDIA DRIVE, and NVIDIA DGX Spark platforms. TensorRT Edge-LLM provides convenient Python scripts to convert HuggingFace checkpoints to [ONNX](https://onnx.ai). Engine build and end-to-end inference runs entirely on Edge platforms. +TensorRT Edge-LLM is NVIDIA's high-performance C++ inference runtime for Large Language Models (LLMs), Vision-Language Models (VLMs), Vision-Language-Action (VLA) models, and World Foundation Models (WFMs) on embedded platforms. It enables efficient deployment of state-of-the-art language models on resource-constrained devices such as NVIDIA Jetson, NVIDIA DRIVE, and NVIDIA DGX Spark platforms. TensorRT Edge-LLM provides convenient Python scripts to convert HuggingFace checkpoints to [ONNX](https://onnx.ai). Engine build and end-to-end inference runs entirely on Edge platforms. --- @@ -25,6 +25,66 @@ For the supported platforms, models and precisions, see the [**Overview**](https --- +## World Foundation Models (Cosmos3-Edge) + +TensorRT Edge-LLM includes `WFMInferenceRuntime` and the `wfm_inference` example for text-to-video generation with [nvidia/Cosmos3-Edge](https://huggingface.co/nvidia/Cosmos3-Edge). The runtime chains five TRT engines (VAE encode/decode, vision embed, MoT backbone, denoise head) behind a single C++ API. + +### Build `wfm_inference` + +After building the C++ runtime and Edge-LLM TensorRT plugins (see [Installation](https://nvidia.github.io/TensorRT-Edge-LLM/latest/user_guide/getting_started/installation.html)): + +```bash +cd /path/to/TensorRT-Edge-LLM/build +cmake --build . --target wfm_inference -j$(nproc) +``` + +### Export engines + +Cosmos3-Edge engines are exported with the companion script `Test/wfm/export_wfm_cosmos_edge.py` (Torch-TensorRT; requires a CUDA GPU with enough memory for the ~3.9B MoT backbone). From the `Test` export workspace: + +```bash +cd /path/to/Test +export EDGE_LLM_PLUGIN_SO=/path/to/TensorRT-Edge-LLM/build-plugin-trt11/libNvInfer_edgellm_plugin.so + +python wfm/export_wfm_cosmos_edge.py \ + --engine-dir /tmp/cosmos_edge_engines \ + --dtype fp16 \ + --num-inference-steps 2 \ + --prompt "A robot arm picks up a red cube." +``` + +This writes an engine bundle: + +``` +/ + config.json + packing_static.json + embedding.safetensors + tokenizer/ + visual_encode/visual_encode.engine + embed/embed.engine + mot_backbone/mot_backbone.engine + denoise_head/denoise_head.engine + visual_decode/visual_decode.engine +``` + +Use the same `--prompt` at export and inference time so `packing_static.json` matches the runtime text packing (`und_len` / sequence length). + +### Run inference + +```bash +export EDGELLM_PLUGIN_PATH=/path/to/TensorRT-Edge-LLM/build-plugin-trt11/libNvInfer_edgellm_plugin.so + +./build/examples/wfm/wfm_inference \ + --engineDir=/tmp/cosmos_edge_engines \ + --inputFile=examples/wfm/wfm_input_example.json \ + --outputFile=/tmp/wfm_output.json +``` + +See `examples/wfm/wfm_input_example.json` for the request JSON format. Each request specifies a text `prompt`, optional `output_video_file` (raw fp16 tensor dump), `num_inference_steps`, and `seed`. Set `generate_sound: true` only for Cosmos3-Omni bundles that include audio engines. + +--- + ## Documentation ### Introduction @@ -79,6 +139,7 @@ See the [**Performance Benchmarks**](https://nvidia.github.io/TensorRT-Edge-LLM/ - Task planning and reasoning - Visual question answering - Human-robot collaboration +- World-model video generation (Cosmos3-Edge WFM) **🏭 Industrial IoT** - Equipment monitoring with NLP diff --git a/cpp/action/actionContextRunner.cpp b/cpp/action/actionContextRunner.cpp new file mode 100644 index 00000000..cdef40f2 --- /dev/null +++ b/cpp/action/actionContextRunner.cpp @@ -0,0 +1,435 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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. + */ + +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "action/actionContextRunner.h" + +#include "common/checkMacros.h" +#include "common/logger.h" +#include "common/mmapReader.h" + +#include +#include +#include +#include + +namespace trt_edgellm +{ +namespace rt +{ +namespace +{ + +//! Builds a Coords from a JSON array of dimension sizes (e.g. [1, 3, 224, 224]). +//! Expects \p shape to be an array of integers; malformed JSON will throw. +rt::Coords coordsFromJson(nlohmann::json const& shape) +{ + std::vector dims; + dims.reserve(shape.size()); + for (auto const& dim : shape) + { + dims.push_back(dim.get()); + } + return rt::Coords(dims); +} + +nvinfer1::DataType dataTypeFromTorchString(std::string const& dtype) +{ + if (dtype == "torch.float16" || dtype == "float16" || dtype == "fp16") + { + return nvinfer1::DataType::kHALF; + } + if (dtype == "torch.float32" || dtype == "float32" || dtype == "fp32") + { + return nvinfer1::DataType::kFLOAT; + } + throw std::runtime_error("ActionContextRunner: unsupported tensor dtype: " + dtype); +} + +std::string resolveIOTensorName(nvinfer1::ICudaEngine const* engine, std::string const& preferredName, + nvinfer1::TensorIOMode mode, int32_t modeIndex = 0) +{ + for (int32_t i = 0; i < engine->getNbIOTensors(); ++i) + { + auto const* name = engine->getIOTensorName(i); + if (name != nullptr && preferredName == name && engine->getTensorIOMode(name) == mode) + { + return preferredName; + } + } + + int32_t seen = 0; + for (int32_t i = 0; i < engine->getNbIOTensors(); ++i) + { + auto const* name = engine->getIOTensorName(i); + if (name != nullptr && engine->getTensorIOMode(name) == mode) + { + if (seen == modeIndex) + { + LOG_WARNING( + "ActionContextRunner: tensor name '%s' is not present in the TensorRT engine; using binding '%s'", + preferredName.c_str(), name); + return name; + } + ++seen; + } + } + + throw std::runtime_error("ActionContextRunner: failed to resolve TensorRT I/O tensor name: " + preferredName); +} + +int64_t tensorBytes(rt::Tensor const& tensor) +{ + return static_cast(tensor.getShape().volume()) * rt::utils::getTypeSize(tensor.getDataType()); +} + +bool copyTensorToDevice(rt::Tensor& dst, rt::Tensor const& src, cudaStream_t stream) +{ + auto const dstBytes = tensorBytes(dst); + auto const srcBytes = tensorBytes(src); + if (dstBytes == srcBytes) + { + CUDA_CHECK(cudaMemcpyAsync(dst.rawPointer(), src.rawPointer(), dstBytes, cudaMemcpyDeviceToDevice, stream)); + return true; + } + + auto const elements = dst.getShape().volume(); + if (src.getShape().volume() != elements) + { + LOG_ERROR("ActionContextRunner: copy shape mismatch, dst=%ld src=%ld", elements, src.getShape().volume()); + return false; + } + + auto const srcType = src.getDataType(); + auto const dstType = dst.getDataType(); + if (srcType == nvinfer1::DataType::kHALF && dstType == nvinfer1::DataType::kFLOAT) + { + std::vector hostSrc(static_cast(elements)); + CUDA_CHECK(cudaMemcpyAsync(hostSrc.data(), src.rawPointer(), srcBytes, cudaMemcpyDeviceToHost, stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + std::vector hostDst(static_cast(elements)); + for (int64_t i = 0; i < elements; ++i) + { + hostDst[static_cast(i)] = __half2float(hostSrc[static_cast(i)]); + } + CUDA_CHECK(cudaMemcpyAsync(dst.rawPointer(), hostDst.data(), dstBytes, cudaMemcpyHostToDevice, stream)); + return true; + } + + if (srcType == nvinfer1::DataType::kFLOAT && dstType == nvinfer1::DataType::kHALF) + { + std::vector hostSrc(static_cast(elements)); + CUDA_CHECK(cudaMemcpyAsync(hostSrc.data(), src.rawPointer(), srcBytes, cudaMemcpyDeviceToHost, stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + std::vector hostDst(static_cast(elements)); + for (int64_t i = 0; i < elements; ++i) + { + hostDst[static_cast(i)] = __float2half(hostSrc[static_cast(i)]); + } + CUDA_CHECK(cudaMemcpyAsync(dst.rawPointer(), hostDst.data(), dstBytes, cudaMemcpyHostToDevice, stream)); + return true; + } + + LOG_ERROR("ActionContextRunner: unsupported dtype conversion src=%d dst=%d", static_cast(srcType), + static_cast(dstType)); + return false; +} + +} // namespace + +ActionContextRunner::ActionContextRunner(std::string const& engineDir, cudaStream_t stream) + : mEngineDir(engineDir) +{ + if (!loadConfig()) + { + throw std::runtime_error("ActionContextRunner: failed to load config"); + } + if (!loadEngine(stream)) + { + throw std::runtime_error("ActionContextRunner: failed to load TensorRT engine"); + } + if (!validateAndFillConfig()) + { + throw std::runtime_error("ActionContextRunner: failed to validate config"); + } + if (!allocateBuffers()) + { + throw std::runtime_error("ActionContextRunner: failed to allocate buffers"); + } +} + +bool ActionContextRunner::loadConfig() +{ + auto const configPath = mEngineDir + "/config.json"; + std::ifstream configFile(configPath); + if (!configFile) + { + LOG_ERROR("ActionContextRunner: failed to open config file: %s", configPath.c_str()); + return false; + } + + try + { + configFile >> mConfigJson; + } + catch (nlohmann::json::parse_error const& e) + { + LOG_ERROR("ActionContextRunner: failed to parse config file: %s", e.what()); + return false; + } + return true; +} + +bool ActionContextRunner::loadEngine(cudaStream_t stream) +{ + auto const engineFile = mConfigJson.value("engine_file", std::string{"action_context.engine"}); + auto const enginePath = mEngineDir + "/" + engineFile; + + file_io::MmapReader engineFileReader(enginePath); + mRuntime.reset(nvinfer1::createInferRuntime(gLogger)); + if (!mRuntime) + { + LOG_ERROR("ActionContextRunner: failed to create TensorRT runtime for %s", enginePath.c_str()); + return false; + } + + mEngine.reset(mRuntime->deserializeCudaEngine(engineFileReader.getData(), engineFileReader.getSize())); + if (!mEngine) + { + LOG_ERROR("ActionContextRunner: failed to deserialize TensorRT engine: %s", enginePath.c_str()); + return false; + } + + mContext.reset(mEngine->createExecutionContext()); + if (!mContext) + { + LOG_ERROR("ActionContextRunner: failed to create TensorRT execution context for %s", enginePath.c_str()); + return false; + } + + if (!mContext->setOptimizationProfileAsync(0, stream)) + { + LOG_ERROR("ActionContextRunner: failed to set optimization profile for %s", enginePath.c_str()); + return false; + } + return true; +} + +int64_t ActionContextRunner::getRequiredContextMemorySize() const +{ + return mEngine ? mEngine->getDeviceMemorySizeV2() : 0; +} + +bool ActionContextRunner::setContextMemory(rt::Tensor& contextMemory) +{ + // action_context uses TensorRT-managed device memory (see loadEngine). + (void) contextMemory; + return true; +} + +bool ActionContextRunner::validateAndFillConfig() +{ + auto const modelTypeStr = mConfigJson.value("model_type", std::string{"action_context"}); + if (modelTypeStr != "action_context") + { + LOG_ERROR("ActionContextRunner: invalid model type: %s", modelTypeStr.c_str()); + return false; + } + + if (mConfigJson.contains("input_names") && !mConfigJson.at("input_names").empty()) + { + mInputName = mConfigJson.at("input_names").at(0).get(); + } + if (mConfigJson.contains("output_names") && !mConfigJson.at("output_names").empty()) + { + mOutputName = mConfigJson.at("output_names").at(0).get(); + } + + if (!mConfigJson.contains("inputs") || !mConfigJson.at("inputs").contains(mInputName)) + { + LOG_ERROR("ActionContextRunner: config is missing input metadata for %s", mInputName.c_str()); + return false; + } + if (!mConfigJson.contains("outputs") || mConfigJson.at("outputs").empty()) + { + LOG_ERROR("ActionContextRunner: config is missing output metadata"); + return false; + } + + auto const& inputMeta = mConfigJson.at("inputs").at(mInputName); + auto const& outputMeta = mConfigJson.at("outputs").at(0); + + mInputShape = coordsFromJson(inputMeta.at("shape")); + mOutputShape = coordsFromJson(outputMeta.at("shape")); + mInputType = dataTypeFromTorchString(inputMeta.at("dtype").get()); + mOutputType = dataTypeFromTorchString(outputMeta.at("dtype").get()); + + if (mInputShape.getNumDims() != 3 || mOutputShape.getNumDims() != 3) + { + LOG_ERROR("ActionContextRunner: expected rank-3 input/output tensors"); + return false; + } + + mMaxBatchSize = static_cast(mInputShape[0]); + mMaxSeqLen = static_cast(mInputShape[1]); + mHiddenSize = static_cast(mInputShape[2]); + mContextHiddenSize = static_cast(mOutputShape[2]); + + mMaxBatchSize = mConfigJson.value("max_batch_size", mMaxBatchSize); + mMaxSeqLen = mConfigJson.value("context_seq_len", mMaxSeqLen); + mContextHiddenSize = mConfigJson.value("context_hidden_size", mContextHiddenSize); + + mInputName = resolveIOTensorName(mEngine.get(), mInputName, nvinfer1::TensorIOMode::kINPUT, 0); + mOutputName = resolveIOTensorName(mEngine.get(), mOutputName, nvinfer1::TensorIOMode::kOUTPUT, 0); + return true; +} + +bool ActionContextRunner::allocateBuffers() +{ + mInputTensor = rt::Tensor(mInputShape, rt::DeviceType::kGPU, mInputType, mInputName); + mOutputTensor = rt::Tensor(mOutputShape, rt::DeviceType::kGPU, mOutputType, mOutputName); + return bindTensors(); +} + +bool ActionContextRunner::reshapeForContext(int32_t batchSize, int32_t seqLen) +{ + if (batchSize <= 0 || batchSize > mMaxBatchSize) + { + LOG_ERROR("ActionContextRunner: batchSize=%d exceeds max batch=%d", batchSize, mMaxBatchSize); + return false; + } + if (seqLen <= 0 || seqLen > mMaxSeqLen) + { + LOG_ERROR("ActionContextRunner: seqLen=%d exceeds max seq=%d", seqLen, mMaxSeqLen); + return false; + } + + auto inputShape = mInputShape; + inputShape[0] = batchSize; + inputShape[1] = seqLen; + auto outputShape = mOutputShape; + outputShape[0] = batchSize; + outputShape[1] = seqLen; + + if (!mInputTensor.reshape(inputShape) || !mOutputTensor.reshape(outputShape)) + { + LOG_ERROR("ActionContextRunner: failed to reshape tensors for batch=%d seq=%d", batchSize, seqLen); + return false; + } + return bindTensors(); +} + +bool ActionContextRunner::bindTensors() noexcept +{ + if (!mContext->setInputShape(mInputName.c_str(), mInputTensor.getShape().getTRTDims())) + { + LOG_ERROR("ActionContextRunner: failed to set input shape for %s", mInputName.c_str()); + return false; + } + if (!mContext->setTensorAddress(mInputName.c_str(), mInputTensor.rawPointer())) + { + LOG_ERROR("ActionContextRunner: failed to bind input tensor %s", mInputName.c_str()); + return false; + } + if (!mContext->setTensorAddress(mOutputName.c_str(), mOutputTensor.rawPointer())) + { + LOG_ERROR("ActionContextRunner: failed to bind output tensor %s", mOutputName.c_str()); + return false; + } + return true; +} + +bool ActionContextRunner::copyLmHiddenFrom(rt::Tensor const& lmHidden, cudaStream_t stream, int32_t validSeqLen) +{ + if (validSeqLen < 0) + { + return copyTensorToDevice(mInputTensor, lmHidden, stream); + } + + if (lmHidden.getShape().getNumDims() != 3 || mInputTensor.getShape().getNumDims() != 3) + { + LOG_ERROR("ActionContextRunner: padded copy expects rank-3 lm_hidden tensors"); + return false; + } + + int32_t const batch = static_cast(mInputTensor.getShape()[0]); + int32_t const dstSeqLen = static_cast(mInputTensor.getShape()[1]); + int32_t const hidden = static_cast(mInputTensor.getShape()[2]); + + CUDA_CHECK(cudaMemsetAsync(mInputTensor.rawPointer(), 0, tensorBytes(mInputTensor), stream)); + if (validSeqLen == 0) + { + return true; + } + + if (validSeqLen > dstSeqLen) + { + LOG_ERROR("ActionContextRunner: invalid validSeqLen=%d for padded copy (dst seq=%d)", validSeqLen, dstSeqLen); + return false; + } + if (lmHidden.getShape()[0] != batch || lmHidden.getShape()[2] != hidden || lmHidden.getShape()[1] < validSeqLen) + { + LOG_ERROR("ActionContextRunner: lmHidden shape %s incompatible with padded copy to %s (validSeqLen=%d)", + lmHidden.getShape().formatString().c_str(), mInputTensor.getShape().formatString().c_str(), validSeqLen); + return false; + } + + size_t const rowBytes = static_cast(hidden) * rt::utils::getTypeSize(mInputTensor.getDataType()); + size_t const copyBytes = static_cast(validSeqLen) * rowBytes; + int64_t const srcSeqLen = lmHidden.getShape()[1]; + for (int32_t b = 0; b < batch; ++b) + { + auto* dstRow = static_cast(mInputTensor.rawPointer()) + static_cast(b) * dstSeqLen * rowBytes; + auto const* srcRow + = static_cast(lmHidden.rawPointer()) + static_cast(b) * srcSeqLen * rowBytes; + CUDA_CHECK(cudaMemcpyAsync(dstRow, srcRow, copyBytes, cudaMemcpyDeviceToDevice, stream)); + } + return true; +} + +bool ActionContextRunner::resetExecutionContext(cudaStream_t stream) +{ + mContext.reset(); + mContext.reset(mEngine->createExecutionContext()); + if (!mContext) + { + LOG_ERROR("ActionContextRunner: failed to recreate TensorRT execution context"); + return false; + } + if (!mContext->setOptimizationProfileAsync(0, stream)) + { + LOG_ERROR("ActionContextRunner: failed to set optimization profile after reset"); + return false; + } + return bindTensors(); +} + +bool ActionContextRunner::infer(cudaStream_t stream) noexcept +{ + if (!bindTensors()) + { + return false; + } + return mContext->enqueueV3(stream); +} + +} // namespace rt +} // namespace trt_edgellm diff --git a/cpp/action/actionContextRunner.h b/cpp/action/actionContextRunner.h new file mode 100644 index 00000000..26da3f9e --- /dev/null +++ b/cpp/action/actionContextRunner.h @@ -0,0 +1,135 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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. + */ + +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "common/tensor.h" + +#include +#include +#include + +#include +#include +#include + +namespace trt_edgellm +{ +namespace rt +{ + +//! GR00T action-context projection engine: lm_hidden_states -> vl_embs. +//! +//! The exported TensorRT engine owns the model-specific stack +//! (eagle_linear -> vlln -> vl_self_attention). The runner only implements the +//! common runtime contract: +//! +//! lm_hidden_states -> vl_embs +class ActionContextRunner +{ +public: + ActionContextRunner(std::string const& engineDir, cudaStream_t stream); + ~ActionContextRunner() noexcept = default; + + int64_t getRequiredContextMemorySize() const; + bool setContextMemory(rt::Tensor& contextMemory); + + int32_t getMaxBatchSize() const noexcept + { + return mMaxBatchSize; + } + int32_t getMaxSeqLen() const noexcept + { + return mMaxSeqLen; + } + int32_t getHiddenSize() const noexcept + { + return mHiddenSize; + } + int32_t getContextHiddenSize() const noexcept + { + return mContextHiddenSize; + } + + std::string const& getInputName() const noexcept + { + return mInputName; + } + std::string const& getOutputName() const noexcept + { + return mOutputName; + } + + rt::Tensor& getLmHiddenStatesInput() noexcept + { + return mInputTensor; + } + rt::Tensor const& getLmHiddenStatesInput() const noexcept + { + return mInputTensor; + } + rt::Tensor& getVlEmbs() noexcept + { + return mOutputTensor; + } + rt::Tensor const& getVlEmbs() const noexcept + { + return mOutputTensor; + } + + bool reshapeForContext(int32_t batchSize, int32_t seqLen); + bool copyLmHiddenFrom(rt::Tensor const& lmHidden, cudaStream_t stream, int32_t validSeqLen = -1); + bool resetExecutionContext(cudaStream_t stream); + bool infer(cudaStream_t stream) noexcept; + +private: + bool loadConfig(); + bool loadEngine(cudaStream_t stream); + bool validateAndFillConfig(); + bool allocateBuffers(); + bool bindTensors() noexcept; + + std::string mEngineDir; + nlohmann::json mConfigJson; + + std::unique_ptr mRuntime; + std::unique_ptr mEngine; + std::unique_ptr mContext; + + rt::Tensor mInputTensor; + rt::Tensor mOutputTensor; + + std::string mInputName{"lm_hidden_states"}; + std::string mOutputName{"vl_embs"}; + rt::Coords mInputShape; + rt::Coords mOutputShape; + nvinfer1::DataType mInputType{nvinfer1::DataType::kHALF}; + nvinfer1::DataType mOutputType{nvinfer1::DataType::kHALF}; + + int32_t mMaxBatchSize{1}; + int32_t mMaxSeqLen{1}; + int32_t mHiddenSize{0}; + int32_t mContextHiddenSize{0}; +}; + +} // namespace rt +} // namespace trt_edgellm diff --git a/cpp/action/actionRunner.cpp b/cpp/action/actionRunner.cpp new file mode 100644 index 00000000..6c008a18 --- /dev/null +++ b/cpp/action/actionRunner.cpp @@ -0,0 +1,1660 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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. + */ + +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "action/actionRunner.h" + +#include "action/actionUtils.h" +#include "common/bindingNames.h" +#include "common/checkMacros.h" +#include "common/logger.h" +#include "common/mmapReader.h" +#include "kernels/posEncoding/initializeCosSinCache.h" + +#include +#include + +#include +#include +#include +#include +#include +#include + +using namespace nvinfer1; +using Json = nlohmann::json; + +namespace trt_edgellm +{ +namespace rt +{ +namespace +{ + +rt::Coords coordsFromJson(Json const& shape) +{ + std::vector dims; + dims.reserve(shape.size()); + for (auto const& dim : shape) + { + dims.push_back(dim.get()); + } + return rt::Coords(dims); +} + +nvinfer1::DataType dataTypeFromTorchString(std::string const& dtype) +{ + if (dtype == "torch.float16" || dtype == "float16" || dtype == "fp16") + { + return nvinfer1::DataType::kHALF; + } + if (dtype == "torch.float32" || dtype == "float32" || dtype == "fp32") + { + return nvinfer1::DataType::kFLOAT; + } + if (dtype == "torch.int32" || dtype == "int32") + { + return nvinfer1::DataType::kINT32; + } + if (dtype == "torch.int64" || dtype == "int64") + { + return nvinfer1::DataType::kINT64; + } + if (dtype == "torch.bool" || dtype == "bool") + { + return nvinfer1::DataType::kBOOL; + } + if (dtype == "torch.uint8" || dtype == "uint8") + { + return nvinfer1::DataType::kUINT8; + } + throw std::runtime_error("Unsupported action tensor dtype: " + dtype); +} + +std::size_t tensorBytes(rt::Tensor const& tensor) +{ + return static_cast(tensor.getShape().volume()) * rt::utils::getTypeSize(tensor.getDataType()); +} + +std::string resolveActionEnginePath(std::string const& engineDir, Json const& configJson) +{ + std::string const engineFile = configJson.value("engine_file", std::string{"action.engine"}); + std::filesystem::path const preferred = std::filesystem::path(engineDir) / engineFile; + if (std::filesystem::exists(preferred)) + { + return preferred.string(); + } + + for (char const* candidate : {"action.engine", "diffusion.engine"}) + { + std::filesystem::path const path = std::filesystem::path(engineDir) / candidate; + if (std::filesystem::exists(path)) + { + return path.string(); + } + } + + return preferred.string(); +} + +std::string resolveVelocityOutputName(Json const& configJson, nvinfer1::ICudaEngine const* engine) +{ + std::string preferred; + if (configJson.contains("output_names") && !configJson.at("output_names").empty()) + { + preferred = configJson.at("output_names").at(0).get(); + } + else if (configJson.contains("outputs") && !configJson.at("outputs").empty() + && configJson.at("outputs").at(0).contains("name")) + { + preferred = configJson.at("outputs").at(0).at("name").get(); + } + else + { + preferred = "velocity"; + } + + for (int32_t i = 0; i < engine->getNbIOTensors(); ++i) + { + char const* name = engine->getIOTensorName(i); + if (name != nullptr && preferred == name && engine->getTensorIOMode(name) == nvinfer1::TensorIOMode::kOUTPUT) + { + return preferred; + } + } + + for (int32_t i = 0; i < engine->getNbIOTensors(); ++i) + { + char const* name = engine->getIOTensorName(i); + if (name != nullptr && engine->getTensorIOMode(name) == nvinfer1::TensorIOMode::kOUTPUT) + { + LOG_WARNING("ActionRunner: output name '%s' not found in engine; using '%s'", preferred.c_str(), name); + return name; + } + } + + throw std::runtime_error("ActionRunner: failed to resolve velocity output tensor name"); +} + +bool bindTensor(nvinfer1::ICudaEngine const* engine, nvinfer1::IExecutionContext* context, std::string const& name, + rt::Tensor& tensor) +{ + if (engine->getTensorIOMode(name.c_str()) == nvinfer1::TensorIOMode::kINPUT) + { + if (!context->setInputShape(name.c_str(), tensor.getShape().getTRTDims())) + { + LOG_ERROR("ActionRunner: failed to set input shape for %s", name.c_str()); + return false; + } + } + + if (!context->setTensorAddress(name.c_str(), tensor.rawPointer())) + { + LOG_ERROR("ActionRunner: failed to bind tensor %s", name.c_str()); + return false; + } + + return true; +} + +bool copyTensorToInput(rt::Tensor& dst, rt::Tensor const& src, cudaStream_t stream) +{ + auto const dstBytes = tensorBytes(dst); + auto const srcBytes = tensorBytes(src); + if (dstBytes == srcBytes) + { + CUDA_CHECK(cudaMemcpyAsync(dst.rawPointer(), src.rawPointer(), dstBytes, cudaMemcpyDeviceToDevice, stream)); + return true; + } + + auto const elements = dst.getShape().volume(); + if (src.getShape().volume() != elements) + { + LOG_ERROR("ActionRunner: copy shape mismatch, dst=%ld src=%ld", elements, src.getShape().volume()); + return false; + } + + auto const srcType = src.getDataType(); + auto const dstType = dst.getDataType(); + if (srcType == nvinfer1::DataType::kHALF && dstType == nvinfer1::DataType::kFLOAT) + { + std::vector hostSrc(static_cast(elements)); + CUDA_CHECK(cudaMemcpyAsync(hostSrc.data(), src.rawPointer(), srcBytes, cudaMemcpyDeviceToHost, stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + std::vector hostDst(static_cast(elements)); + for (int64_t i = 0; i < elements; ++i) + { + hostDst[static_cast(i)] = __half2float(hostSrc[static_cast(i)]); + } + CUDA_CHECK(cudaMemcpyAsync(dst.rawPointer(), hostDst.data(), dstBytes, cudaMemcpyHostToDevice, stream)); + return true; + } + + if (srcType == nvinfer1::DataType::kFLOAT && dstType == nvinfer1::DataType::kHALF) + { + std::vector hostSrc(static_cast(elements)); + CUDA_CHECK(cudaMemcpyAsync(hostSrc.data(), src.rawPointer(), srcBytes, cudaMemcpyDeviceToHost, stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + std::vector hostDst(static_cast(elements)); + for (int64_t i = 0; i < elements; ++i) + { + hostDst[static_cast(i)] = __float2half(hostSrc[static_cast(i)]); + } + CUDA_CHECK(cudaMemcpyAsync(dst.rawPointer(), hostDst.data(), dstBytes, cudaMemcpyHostToDevice, stream)); + return true; + } + + LOG_ERROR("ActionRunner: unsupported dtype conversion src=%d dst=%d", static_cast(srcType), + static_cast(dstType)); + return false; +} + +void fillNormal(rt::Tensor& tensor, std::mt19937& rng) +{ + std::normal_distribution distribution(0.0F, 1.0F); + auto const elements = tensor.getShape().volume(); + + if (tensor.getDataType() == nvinfer1::DataType::kFLOAT) + { + auto* data = tensor.dataPointer(); + for (int64_t i = 0; i < elements; ++i) + { + data[i] = distribution(rng); + } + return; + } + + if (tensor.getDataType() == nvinfer1::DataType::kHALF) + { + auto* data = tensor.dataPointer(); + for (int64_t i = 0; i < elements; ++i) + { + data[i] = __float2half(distribution(rng)); + } + return; + } + + throw std::runtime_error("ActionRunner: noise initialization only supports fp16/fp32 actions"); +} + +void updateActionTensor(rt::Tensor& actions, rt::Tensor const& predVelocity, float stepSize) +{ + auto const elements = actions.getShape().volume(); + check::check(actions.getDataType() == predVelocity.getDataType(), "ActionRunner dtype mismatch"); + check::check(actions.getShape().volume() == predVelocity.getShape().volume(), "ActionRunner shape mismatch"); + + if (actions.getDataType() == nvinfer1::DataType::kFLOAT) + { + auto* actionData = actions.dataPointer(); + auto const* predData = predVelocity.dataPointer(); + for (int64_t i = 0; i < elements; ++i) + { + actionData[i] += stepSize * predData[i]; + } + return; + } + + if (actions.getDataType() == nvinfer1::DataType::kHALF) + { + auto* actionData = actions.dataPointer(); + auto const* predData = predVelocity.dataPointer(); + for (int64_t i = 0; i < elements; ++i) + { + auto const updated = __half2float(actionData[i]) + stepSize * __half2float(predData[i]); + actionData[i] = __float2half(updated); + } + return; + } + + throw std::runtime_error("ActionRunner: action update only supports fp16/fp32 actions"); +} + +bool fillHostTensorFromFloats(rt::Tensor& tensor, std::vector const& values) +{ + auto const elements = tensor.getShape().volume(); + if (static_cast(elements) != values.size()) + { + LOG_ERROR("ActionRunner: expected %ld values but got %zu", elements, values.size()); + return false; + } + + if (tensor.getDataType() == nvinfer1::DataType::kFLOAT) + { + auto* data = tensor.dataPointer(); + for (int64_t i = 0; i < elements; ++i) + { + data[i] = values[static_cast(i)]; + } + return true; + } + + if (tensor.getDataType() == nvinfer1::DataType::kHALF) + { + auto* data = tensor.dataPointer(); + for (int64_t i = 0; i < elements; ++i) + { + data[i] = __float2half(values[static_cast(i)]); + } + return true; + } + + if (tensor.getDataType() == nvinfer1::DataType::kINT64) + { + auto* data = tensor.dataPointer(); + for (int64_t i = 0; i < elements; ++i) + { + data[i] = static_cast(values[static_cast(i)]); + } + return true; + } + + if (tensor.getDataType() == nvinfer1::DataType::kINT32) + { + auto* data = tensor.dataPointer(); + for (int64_t i = 0; i < elements; ++i) + { + data[i] = static_cast(values[static_cast(i)]); + } + return true; + } + + LOG_ERROR("ActionRunner: unsupported tensor dtype for host fill: %d", static_cast(tensor.getDataType())); + return false; +} + +} // namespace + +ActionRunner::ActionRunner( + std::string const& engineDir, cudaStream_t stream, LinearKVCache::CacheConfig const& kvCacheConfig) + : mStream(stream) +{ + LOG_DEBUG("Loading action runner from %s", engineDir.c_str()); + + std::string const configPath = engineDir + "/config.json"; + { + std::ifstream configFile(configPath); + if (configFile) + { + try + { + configFile >> mConfigJson; + } + catch (Json::parse_error const& e) + { + LOG_WARNING("ActionRunner: failed to parse %s: %s", configPath.c_str(), e.what()); + mConfigJson = Json::object(); + } + } + } + + std::string const actionEnginePath = resolveActionEnginePath(engineDir, mConfigJson); + + mRuntime = std::unique_ptr(createInferRuntime(gLogger)); + if (!mRuntime) + { + throw std::runtime_error("Failed to create TensorRT runtime"); + } + + auto mmapReader = std::make_unique(actionEnginePath); + if (mmapReader->getData() == nullptr) + { + throw std::runtime_error("Failed to read engine file: " + actionEnginePath); + } + + mEngine + = std::unique_ptr(mRuntime->deserializeCudaEngine(mmapReader->getData(), mmapReader->getSize())); + if (!mEngine) + { + throw std::runtime_error("Failed to deserialize engine from: " + actionEnginePath); + } + + mContext = std::unique_ptr( + mEngine->createExecutionContext(ExecutionContextAllocationStrategy::kUSER_MANAGED)); + if (!mContext) + { + throw std::runtime_error("Failed to create execution context"); + } + + if (!mContext->setOptimizationProfileAsync(0, stream)) + { + throw std::runtime_error("Failed to set optimization profile"); + } + + if (!parseModelConfig(configPath)) + { + throw std::runtime_error("Failed to parse model config"); + } + + try + { + allocateTensors(kvCacheConfig); + } + catch (std::exception const& e) + { + LOG_ERROR("ActionRunner tensor allocation failed: %s", e.what()); + throw; + } + + CUDA_CHECK(cudaStreamSynchronize(stream)); +} + +int64_t ActionRunner::getRequiredContextMemorySize() const +{ + return mEngine ? mEngine->getDeviceMemorySizeV2() : 0; +} + +bool ActionRunner::setContextMemory(rt::Tensor& sharedContextMemory) +{ + int64_t const requiredSize = getRequiredContextMemorySize(); + if (sharedContextMemory.getMemoryCapacity() < requiredSize) + { + LOG_ERROR("Shared context memory (%lld bytes) is smaller than required (%lld bytes)", + static_cast(sharedContextMemory.getMemoryCapacity()), static_cast(requiredSize)); + return false; + } + mExecContextMemory = sharedContextMemory.rawPointer(); + mExecContextMemoryCapacity = sharedContextMemory.getMemoryCapacity(); + mContext->setDeviceMemoryV2(mExecContextMemory, mExecContextMemoryCapacity); + return true; +} + +bool ActionRunner::resetExecutionContext(cudaStream_t stream) +{ + if (!mEngine || mExecContextMemory == nullptr || mExecContextMemoryCapacity <= 0) + { + LOG_ERROR("ActionRunner: cannot reset execution context before shared memory is configured"); + return false; + } + + mContext.reset(); + mContext.reset(mEngine->createExecutionContext(ExecutionContextAllocationStrategy::kUSER_MANAGED)); + if (!mContext) + { + LOG_ERROR("ActionRunner: failed to recreate execution context"); + return false; + } + if (!mContext->setOptimizationProfileAsync(0, stream)) + { + LOG_ERROR("ActionRunner: failed to set optimization profile after reset"); + return false; + } + mContext->setDeviceMemoryV2(mExecContextMemory, mExecContextMemoryCapacity); + return true; +} + +bool ActionRunner::engineHasTensor(std::string const& name) const noexcept +{ + if (!mEngine) + { + return false; + } + for (int32_t i = 0; i < mEngine->getNbIOTensors(); ++i) + { + char const* tensorName = mEngine->getIOTensorName(i); + if (tensorName != nullptr && name == tensorName) + { + return true; + } + } + return false; +} + +bool ActionRunner::parseModelConfig(std::string const& configPath) +{ + if (mConfigJson.empty()) + { + std::ifstream configFile(configPath); + if (!configFile.is_open()) + { + LOG_ERROR("Failed to open config file: %s", configPath.c_str()); + return false; + } + try + { + configFile >> mConfigJson; + } + catch (Json::parse_error const& e) + { + LOG_ERROR("Failed to parse config file %s: %s", configPath.c_str(), e.what()); + return false; + } + } + + std::string const modelTypeStr = mConfigJson.value("model_type", std::string{}); + if (modelTypeStr == "alpamayo1" || modelTypeStr == "ALPAMAYO1") + { + mModelType = action::ActionModelType::ALPAMAYO1; + } + + bool const hasPrefixKV + = engineHasTensor(binding_names::kNoiseTrajectory) && engineHasTensor(binding_names::formatKCacheName(0, true)); + bool const hasVelocityInputs = mConfigJson.contains("input_names") && mConfigJson.contains("inputs"); + + if (hasPrefixKV || mModelType == action::ActionModelType::ALPAMAYO1) + { + mContextHandoff = ActionContextHandoff::PREFIX_KV; + mRolloutMode = ActionRolloutMode::FLOW_MATCHING; + mModelType = action::ActionModelType::ALPAMAYO1; + mUsesMRope = engineHasTensor(binding_names::kRopeCosSin); + + try + { + mConfig.ropeTheta = mConfigJson.at("rope_theta").get(); + mConfig.numDecoderLayers = mConfigJson.at("num_hidden_layers").get(); + mConfig.numTrajTokens = mConfigJson.at("num_traj_tokens").get(); + mConfig.trajTokenStart = mConfigJson.at("traj_token_start").get(); + } + catch (Json::exception const& e) + { + LOG_ERROR("Failed to read Alpamayo fields from %s: %s", configPath.c_str(), e.what()); + return false; + } + + auto const& ropeParams = mConfigJson.contains("rope_scaling") ? mConfigJson["rope_scaling"] : mConfigJson; + if (ropeParams.contains("mrope_section")) + { + auto section = ropeParams["mrope_section"].get>(); + if (section.size() >= 3) + { + mConfig.mropeSectionH = section[1]; + mConfig.mropeSectionW = section[2]; + } + } + if (mUsesMRope && (mConfig.mropeSectionH <= 0 || mConfig.mropeSectionW <= 0)) + { + LOG_ERROR("ActionRunner: mrope_section required when rope_cos_sin binding is present"); + return false; + } + + if (mConfigJson.contains("builder_config") && mConfigJson["builder_config"].contains("max_kv_cache_capacity")) + { + mConfig.maxKVCacheCapacity = mConfigJson["builder_config"]["max_kv_cache_capacity"].get(); + } + else + { + LOG_ERROR("max_kv_cache_capacity not found in builder_config section of %s", configPath.c_str()); + return false; + } + + mConfig.numInferenceTimesteps = mConfigJson.value("num_inference_timesteps", kDefaultDenoiseSteps); + if (mConfig.numInferenceTimesteps <= 0) + { + mConfig.numInferenceTimesteps = kDefaultDenoiseSteps; + } + + Dims const maxNoise = mEngine->getProfileShape(binding_names::kNoiseTrajectory, 0, OptProfileSelector::kMAX); + mConfig.actionHorizon = static_cast(maxNoise.d[1]); + mConfig.actionDim = static_cast(maxNoise.d[2]); + return true; + } + + if (hasVelocityInputs) + { + mContextHandoff = ActionContextHandoff::CONTEXT_TENSOR; + mRolloutMode = ActionRolloutMode::VELOCITY; + + mConfig.numInferenceTimesteps = mConfigJson.value("num_inference_timesteps", 0); + if (mConfig.numInferenceTimesteps <= 0) + { + mConfig.numInferenceTimesteps = mConfigJson.value("num_inference_steps", 1); + } + mNumTimestepBuckets = mConfigJson.value("num_timestep_buckets", 1); + mRolloutDtSign = mConfigJson.value("rollout_dt_sign", 1); + + mNoiseInputName = mConfigJson.value("noise_input_name", std::string{}); + if (mNoiseInputName.empty()) + { + if (engineHasTensor("actions")) + { + mNoiseInputName = "actions"; + } + else if (engineHasTensor("x_t")) + { + mNoiseInputName = "x_t"; + } + else if (engineHasTensor(binding_names::kNoiseTrajectory)) + { + mNoiseInputName = binding_names::kNoiseTrajectory; + } + else + { + LOG_ERROR("ActionRunner: could not resolve noise input name"); + return false; + } + } + + mTimestepSchedule = mConfigJson.value("timestep_schedule", std::string{"discrete_buckets"}); + if (mTimestepSchedule == "groot_buckets") + { + mTimestepSchedule = "discrete_buckets"; + } + else if (mTimestepSchedule == "pi05_flow") + { + mTimestepSchedule = "continuous_flow"; + } + + mInputNames.clear(); + for (auto const& nameJson : mConfigJson.at("input_names")) + { + auto const name = nameJson.get(); + if (!mConfigJson.at("inputs").contains(name)) + { + LOG_ERROR("ActionRunner: missing input metadata for %s", name.c_str()); + return false; + } + mInputNames.push_back(name); + } + + if (!hasInputTensor(mNoiseInputName) || !hasInputTensor(mTimestepName)) + { + LOG_ERROR("ActionRunner: velocity config must define '%s' and '%s'", mNoiseInputName.c_str(), + mTimestepName.c_str()); + return false; + } + + buildEngineBindingNames(); + + mLmToActionSlots.clear(); + mLmWiredInputNames.clear(); + if (mConfigJson.contains("lm_to_action_slots")) + { + for (auto const& pairJson : mConfigJson.at("lm_to_action_slots")) + { + if (!pairJson.is_array() || pairJson.size() != 2) + { + LOG_ERROR("ActionRunner: lm_to_action_slots entries must be [lm_idx, action_idx]"); + return false; + } + auto const lmSlot = pairJson.at(0).get(); + auto const actionSlot = pairJson.at(1).get(); + mLmToActionSlots.emplace_back(lmSlot, actionSlot); + mLmWiredInputNames.push_back(mInputNames[static_cast(actionSlot)]); + } + } + + auto const actionShape = coordsFromJson(mConfigJson.at("inputs").at(mNoiseInputName).at("shape")); + if (actionShape.getNumDims() > 0) + { + mMaxActionBatchSize = static_cast(actionShape[0]); + if (actionShape.getNumDims() > 1) + { + mConfig.actionHorizon = static_cast(actionShape[1]); + } + if (actionShape.getNumDims() > 2) + { + mConfig.actionDim = static_cast(actionShape[2]); + } + } + + mPredVelocityName = resolveVelocityOutputName(mConfigJson, mEngine.get()); + return true; + } + + LOG_ERROR("ActionRunner: unrecognized action engine layout in %s", configPath.c_str()); + return false; +} + +void ActionRunner::allocateTensors(LinearKVCache::CacheConfig const& kvCacheConfig) +{ + if (mContextHandoff == ActionContextHandoff::PREFIX_KV) + { + allocatePrefixKVTensors(kvCacheConfig); + return; + } + if (mContextHandoff == ActionContextHandoff::CONTEXT_TENSOR) + { + allocateVelocityTensors(); + return; + } + throw std::runtime_error("ActionRunner: cannot allocate tensors for unknown context handoff"); +} + +void ActionRunner::allocatePrefixKVTensors(LinearKVCache::CacheConfig const& kvCacheConfig) +{ + Dims const maxNoise = mEngine->getProfileShape(binding_names::kNoiseTrajectory, 0, OptProfileSelector::kMAX); + int32_t const maxBatch = static_cast(maxNoise.d[0]); + mMaxActionBatchSize = maxBatch; + + mNumKVHeads = static_cast(kvCacheConfig.numKVHeads); + mMaxSequenceLength = static_cast(kvCacheConfig.maxSequenceLength); + mKvHeadDim = static_cast(kvCacheConfig.headDim); + + nvinfer1::Dims noiseShape = mContext->getTensorShape(binding_names::kNoiseTrajectory); + mConfig.actionHorizon = static_cast(noiseShape.d[1]); + mConfig.actionDim = static_cast(noiseShape.d[2]); + + rt::Coords const noiseCoords({maxBatch, mConfig.actionHorizon, mConfig.actionDim}); + + mNoiseDevice + = rt::Tensor(noiseCoords, rt::DeviceType::kGPU, nvinfer1::DataType::kFLOAT, "ActionRunner::mNoiseDevice"); + mNoiseHost = rt::Tensor(noiseCoords, rt::DeviceType::kCPU, nvinfer1::DataType::kFLOAT, "ActionRunner::mNoiseHost"); + mDenoisedDevice + = rt::Tensor(noiseCoords, rt::DeviceType::kGPU, nvinfer1::DataType::kFLOAT, "ActionRunner::mDenoisedDevice"); + mDenoisedHost + = rt::Tensor(noiseCoords, rt::DeviceType::kCPU, nvinfer1::DataType::kFLOAT, "ActionRunner::mDenoisedHost"); + + mTimeStepsT0Device + = rt::Tensor({1}, rt::DeviceType::kGPU, nvinfer1::DataType::kFLOAT, "ActionRunner::mTimeStepsT0Device"); + mTimeStepsT1Device + = rt::Tensor({1}, rt::DeviceType::kGPU, nvinfer1::DataType::kFLOAT, "ActionRunner::mTimeStepsT1Device"); + mTimeStepsT0Host = rt::Tensor(std::vector{mConfig.numInferenceTimesteps}, rt::DeviceType::kCPU, + nvinfer1::DataType::kFLOAT, "ActionRunner::mTimeStepsT0Host"); + mTimeStepsT1Host = rt::Tensor(std::vector{mConfig.numInferenceTimesteps}, rt::DeviceType::kCPU, + nvinfer1::DataType::kFLOAT, "ActionRunner::mTimeStepsT1Host"); + + if (mUsesMRope) + { + nvinfer1::Dims ropeCosSinDims = mContext->getTensorShape(binding_names::kRopeCosSin); + mRopeHeadDim = ropeCosSinDims.d[2]; + mRopeCosSinDevice = rt::Tensor({maxBatch, mConfig.actionHorizon, mRopeHeadDim}, rt::DeviceType::kGPU, + nvinfer1::DataType::kFLOAT, "ActionRunner::mRopeCosSinDevice"); + mRopePositionIdsHost = rt::Tensor({maxBatch, 3, mConfig.actionHorizon}, rt::DeviceType::kCPU, + nvinfer1::DataType::kINT64, "ActionRunner::mRopePositionIdsHost"); + mRopePositionIdsDevice = rt::Tensor({maxBatch, 3, mConfig.actionHorizon}, rt::DeviceType::kGPU, + nvinfer1::DataType::kINT64, "ActionRunner::mRopePositionIdsDevice"); + } + + mPositionIdsHost = rt::Tensor({maxBatch, mConfig.actionHorizon}, rt::DeviceType::kCPU, nvinfer1::DataType::kINT32, + "ActionRunner::mPositionIdsHost"); + mPositionIdsDevice = rt::Tensor({maxBatch, mConfig.actionHorizon}, rt::DeviceType::kGPU, nvinfer1::DataType::kINT32, + "ActionRunner::mPositionIdsDevice"); + + mKvcacheActualLengthsHost = rt::Tensor(std::vector{maxBatch}, rt::DeviceType::kCPU, + nvinfer1::DataType::kINT32, "ActionRunner::mKvcacheActualLengthsHost"); + mKvcacheActualLengthsBroadcastDevice = rt::Tensor(std::vector{maxBatch}, rt::DeviceType::kGPU, + nvinfer1::DataType::kINT32, "ActionRunner::mKvcacheActualLengthsBroadcastDevice"); + + rt::Coords const kvcacheShape{ + maxBatch, kvCacheConfig.numKVHeads, kvCacheConfig.maxSequenceLength, kvCacheConfig.headDim}; + mKCacheLayers.resize(static_cast(mConfig.numDecoderLayers)); + mVCacheLayers.resize(static_cast(mConfig.numDecoderLayers)); + for (int32_t i = 0; i < mConfig.numDecoderLayers; ++i) + { + mKCacheLayers[i] = rt::Tensor( + kvcacheShape, rt::DeviceType::kGPU, kvCacheConfig.kvCacheTypeTRT, "ActionRunner::mKCacheLayer"); + mVCacheLayers[i] = rt::Tensor( + kvcacheShape, rt::DeviceType::kGPU, kvCacheConfig.kvCacheTypeTRT, "ActionRunner::mVCacheLayer"); + } +} + +void ActionRunner::allocateVelocityTensors() +{ + mInputTensors.clear(); + mInputTensors.reserve(mInputNames.size()); + for (auto const& name : mInputNames) + { + mInputTensors.emplace_back(coordsFromJson(mConfigJson.at("inputs").at(name).at("shape")), rt::DeviceType::kGPU, + dataTypeFromTorchString(mConfigJson.at("inputs").at(name).at("dtype").get()), name); + } + + mPredVelocity = rt::Tensor(coordsFromJson(mConfigJson.at("outputs").at(0).at("shape")), rt::DeviceType::kGPU, + dataTypeFromTorchString(mConfigJson.at("outputs").at(0).at("dtype").get()), mPredVelocityName); + + mNoiseHost = rt::Tensor( + getActions().getShape(), rt::DeviceType::kCPU, getActions().getDataType(), "ActionRunner::mNoiseHost"); + mPredVelocityHost = rt::Tensor( + mPredVelocity.getShape(), rt::DeviceType::kCPU, mPredVelocity.getDataType(), "ActionRunner::mPredVelocityHost"); + + auto const& timestepMeta = mConfigJson.at("inputs").at(mTimestepName); + mTimestepHost = rt::Tensor(coordsFromJson(timestepMeta.at("shape")), rt::DeviceType::kCPU, + dataTypeFromTorchString(timestepMeta.at("dtype").get()), "ActionRunner::mTimestepHost"); + + if (!bindVelocityTensors()) + { + throw std::runtime_error("ActionRunner: failed to bind velocity tensors during allocation"); + } +} + +bool ActionRunner::reshapeActionTensorsForActiveBatch(int32_t activeBatchSize) +{ + if (mContextHandoff == ActionContextHandoff::PREFIX_KV) + { + rt::Coords const noiseShape({activeBatchSize, mConfig.actionHorizon, mConfig.actionDim}); + bool ok = true; + ok &= mNoiseDevice.reshape(noiseShape); + ok &= mNoiseHost.reshape(noiseShape); + ok &= mDenoisedDevice.reshape(noiseShape); + ok &= mDenoisedHost.reshape(noiseShape); + if (mUsesMRope) + { + ok &= mRopeCosSinDevice.reshape({activeBatchSize, mConfig.actionHorizon, mRopeHeadDim}); + ok &= mRopePositionIdsHost.reshape({activeBatchSize, 3, mConfig.actionHorizon}); + ok &= mRopePositionIdsDevice.reshape({activeBatchSize, 3, mConfig.actionHorizon}); + } + ok &= mPositionIdsHost.reshape({activeBatchSize, mConfig.actionHorizon}); + ok &= mPositionIdsDevice.reshape({activeBatchSize, mConfig.actionHorizon}); + ok &= mKvcacheActualLengthsHost.reshape({activeBatchSize}); + if (!ok) + { + LOG_ERROR("ActionRunner: tensor reshape failed for activeBatchSize=%d", activeBatchSize); + } + return ok; + } + + if (mContextHandoff == ActionContextHandoff::CONTEXT_TENSOR) + { + bool ok = true; + ok &= mNoiseHost.reshape({activeBatchSize, mConfig.actionHorizon, mConfig.actionDim}); + for (std::size_t inputIdx = 0; inputIdx < mInputTensors.size(); ++inputIdx) + { + auto const& inputName = mInputNames[inputIdx]; + if (isRolloutManagedInput(inputName)) + { + continue; + } + auto& tensor = mInputTensors[inputIdx]; + auto shape = tensor.getShape(); + if (shape.getNumDims() > 0) + { + shape[0] = activeBatchSize; + ok &= tensor.reshape(shape); + } + } + auto predShape = mPredVelocity.getShape(); + if (predShape.getNumDims() > 0) + { + predShape[0] = activeBatchSize; + ok &= mPredVelocity.reshape(predShape); + ok &= mPredVelocityHost.reshape(predShape); + } + if (!ok) + { + LOG_ERROR("ActionRunner: velocity tensor reshape failed for activeBatchSize=%d", activeBatchSize); + } + return ok; + } + + return false; +} + +void ActionRunner::initializeNoiseTrajectory(int32_t randomSeed, int32_t activeBatchSize) +{ + if (mContextHandoff == ActionContextHandoff::PREFIX_KV) + { + size_t const elemCount = static_cast(activeBatchSize) * static_cast(mConfig.actionHorizon) + * static_cast(mConfig.actionDim); + float* data = static_cast(mNoiseHost.rawPointer()); + std::mt19937 generator(static_cast(randomSeed)); + std::normal_distribution dist(0.0f, 1.0f); + for (size_t i = 0; i < elemCount; ++i) + { + data[i] = dist(generator); + } + return; + } + + if (mContextHandoff == ActionContextHandoff::CONTEXT_TENSOR) + { + mRng.seed(static_cast(randomSeed)); + fillNormal(mNoiseHost, mRng); + } +} + +bool ActionRunner::preprocess(LLMGenerationRequest const& request, std::vector>& batchedInputIds, + tokenizer::Tokenizer const* tokenizer) +{ + int32_t const numUniqueRequests = static_cast(request.requests.size()); + int32_t const actionBatchSize = (request.actionBatchSize > 0) ? request.actionBatchSize : numUniqueRequests; + + if (mContextHandoff == ActionContextHandoff::PREFIX_KV && mConfig.numTrajTokens > 0 && tokenizer != nullptr) + { + tokenizer::TokenToRanks const& specialTokens = tokenizer->getSpecialTokensEncoder(); + for (int32_t i = 0; i < numUniqueRequests; ++i) + { + std::vector& tokenIds = batchedInputIds[i]; + LLMGenerationRequest::Request const& req = request.requests[i]; + if (!req.pastTrajectory) + { + continue; + } + + auto const itStart = specialTokens.find(kTrajHistoryStartStr); + auto const itPad = specialTokens.find(kTrajHistoryPadStr); + auto const itEnd = specialTokens.find(kTrajHistoryEndStr); + if (itStart == specialTokens.end() || itPad == specialTokens.end() || itEnd == specialTokens.end()) + { + continue; + } + + tokenizer::Rank const startId = itStart->second; + tokenizer::Rank const padId = itPad->second; + tokenizer::Rank const endId = itEnd->second; + + std::vector const actualTokens = action_utils::trajectoryToTokenIds( + *req.pastTrajectory, mConfig.numTrajTokens, mConfig.trajTokenStart); + + size_t scanIdx = 0; + while (scanIdx < tokenIds.size()) + { + if (tokenIds[scanIdx] != startId) + { + ++scanIdx; + continue; + } + size_t endMarkerIdx = scanIdx + 1; + while (endMarkerIdx < tokenIds.size() && tokenIds[endMarkerIdx] == padId) + { + ++endMarkerIdx; + } + if (endMarkerIdx >= tokenIds.size() || tokenIds[endMarkerIdx] != endId) + { + ++scanIdx; + continue; + } + + size_t const numPads = endMarkerIdx - scanIdx - 1; + if (numPads != actualTokens.size()) + { + LOG_ERROR("Trajectory placeholder token count (%zu) does not match encoding length (%zu).", numPads, + actualTokens.size()); + return false; + } + + for (size_t k = 0; k < actualTokens.size(); ++k) + { + tokenIds[scanIdx + 1 + k] = actualTokens[k]; + } + scanIdx += 2 + actualTokens.size(); + } + } + } + + if (actionBatchSize > mMaxActionBatchSize) + { + LOG_ERROR( + "Requested action batch size %d exceeds engine max batch size %d", actionBatchSize, mMaxActionBatchSize); + return false; + } + + mActiveActionBatchSize = actionBatchSize; + if (mContextHandoff == ActionContextHandoff::CONTEXT_TENSOR && !reshapeActionTensorsForActiveBatch(actionBatchSize)) + { + LOG_ERROR("ActionRunner: failed to reshape tensors for action batch size %d", actionBatchSize); + return false; + } + + try + { + initializeNoiseTrajectory(mNoiseSeed, actionBatchSize); + } + catch (std::exception const& e) + { + LOG_ERROR("ActionRunner noise initialization failed: %s", e.what()); + return false; + } + + return true; +} + +std::pair ActionRunner::getSeparateKVCacheForDecoderLayer( + cudaStream_t stream, LinearKVCache& kvcache, int32_t decoderLayerIdx, int32_t activeBatchSize) +{ + LinearKVCache::CacheConfig const& config = kvcache.getConfig(); + int64_t const blockElems = config.numKVHeads * config.maxSequenceLength * config.headDim; + size_t const elemSize = rt::utils::getTypeSize(config.kvCacheTypeTRT); + size_t const blockBytes = static_cast(blockElems) * elemSize; + size_t const combinedBatchStride = 2ULL * blockBytes; + + rt::Tensor combined = kvcache.getCombinedKVCacheForDecoderLayer(decoderLayerIdx); + char const* src = static_cast(combined.rawPointer()); + char* dstK = static_cast(mKCacheLayers[decoderLayerIdx].rawPointer()); + char* dstV = static_cast(mVCacheLayers[decoderLayerIdx].rawPointer()); + + int32_t const llmBatch = static_cast(config.maxBatchSize); + for (int32_t b = 0; b < activeBatchSize; ++b) + { + int32_t const srcSlot = b % llmBatch; + CUDA_CHECK(cudaMemcpyAsync(dstK + static_cast(b) * blockBytes, + src + static_cast(srcSlot) * combinedBatchStride, blockBytes, cudaMemcpyDeviceToDevice, stream)); + CUDA_CHECK(cudaMemcpyAsync(dstV + static_cast(b) * blockBytes, + src + static_cast(srcSlot) * combinedBatchStride + blockBytes, blockBytes, cudaMemcpyDeviceToDevice, + stream)); + } + return {mKCacheLayers[decoderLayerIdx], mVCacheLayers[decoderLayerIdx]}; +} + +void ActionRunner::setDynamicInputShapes(int32_t activeBatchSize) +{ + if (mContextHandoff != ActionContextHandoff::PREFIX_KV) + { + return; + } + + Dims const kvCacheStartIndexShape = {1, {activeBatchSize}}; + Dims const noiseShape = {3, {activeBatchSize, mConfig.actionHorizon, mConfig.actionDim}}; + Dims const ropeIdxShape = {2, {activeBatchSize, mConfig.actionHorizon}}; + Dims const kvShape = { + 4, {activeBatchSize, static_cast(mNumKVHeads), static_cast(mMaxSequenceLength), mKvHeadDim}}; + + bool status = true; + status &= mContext->setInputShape(binding_names::kKVCacheStartIndex, kvCacheStartIndexShape); + status &= mContext->setInputShape(binding_names::kNoiseTrajectory, noiseShape); + if (mUsesMRope) + { + Dims const ropeCosShape = {3, {activeBatchSize, mConfig.actionHorizon, mRopeHeadDim}}; + status &= mContext->setInputShape(binding_names::kRopeCosSin, ropeCosShape); + } + status &= mContext->setInputShape(binding_names::kAttentionPosId, ropeIdxShape); + for (int32_t i = 0; i < mConfig.numDecoderLayers; ++i) + { + status &= mContext->setInputShape(binding_names::formatKCacheName(i, true).c_str(), kvShape); + status &= mContext->setInputShape(binding_names::formatVCacheName(i, true).c_str(), kvShape); + } + if (!status) + { + LOG_ERROR("ActionRunner: setDynamicInputShapes failed for activeBatchSize=%d", activeBatchSize); + throw std::runtime_error("setDynamicInputShapes failed"); + } +} + +int32_t const* ActionRunner::getActualKVLengths(cudaStream_t stream, int32_t activeBatchSize) +{ + CUDA_CHECK(cudaMemcpyAsync(mKvcacheActualLengthsHost.rawPointer(), mKvcacheActualLengthsDevice, + static_cast(activeBatchSize) * sizeof(int32_t), cudaMemcpyDeviceToHost, stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + return static_cast(mKvcacheActualLengthsHost.rawPointer()); +} + +std::vector> ActionRunner::sampleTrajectory(cudaStream_t stream, + int32_t activeBatchSize, LinearKVCache& kvcache, std::vector const& vlmOutputsRopeDeltas) +{ + std::vector> result; + if (mContextHandoff != ActionContextHandoff::PREFIX_KV) + { + LOG_ERROR("ActionRunner::sampleTrajectory requires PREFIX_KV context handoff"); + return result; + } + + if (mUsesMRope && static_cast(vlmOutputsRopeDeltas.size()) != activeBatchSize) + { + LOG_ERROR("vlmOutputsRopeDeltas size %zu != activeBatchSize %d", vlmOutputsRopeDeltas.size(), activeBatchSize); + return result; + } + + if (!reshapeActionTensorsForActiveBatch(activeBatchSize)) + { + LOG_ERROR("reshapeActionTensorsForActiveBatch failed for activeBatchSize=%d", activeBatchSize); + return result; + } + + int32_t const llmBatch = static_cast(kvcache.getConfig().maxBatchSize); + auto* llmLengthsDevice = static_cast(kvcache.getKVCacheLengths().rawPointer()); + if (llmBatch < activeBatchSize) + { + auto* broadcastLengths = static_cast(mKvcacheActualLengthsBroadcastDevice.rawPointer()); + for (int32_t b = 0; b < activeBatchSize; ++b) + { + int32_t const srcSlot = b % llmBatch; + CUDA_CHECK(cudaMemcpyAsync( + broadcastLengths + b, llmLengthsDevice + srcSlot, sizeof(int32_t), cudaMemcpyDeviceToDevice, stream)); + } + mKvcacheActualLengthsDevice = broadcastLengths; + } + else + { + mKvcacheActualLengthsDevice = llmLengthsDevice; + } + + setDynamicInputShapes(activeBatchSize); + + bool setEngineIOStatus{true}; + setEngineIOStatus &= mContext->setTensorAddress(binding_names::kKVCacheStartIndex, mKvcacheActualLengthsDevice); + setEngineIOStatus &= mContext->setTensorAddress(binding_names::kNoiseTrajectory, mNoiseDevice.rawPointer()); + setEngineIOStatus &= mContext->setTensorAddress(binding_names::kTimeStepsT0, mTimeStepsT0Device.rawPointer()); + setEngineIOStatus &= mContext->setTensorAddress(binding_names::kTimeStepsT1, mTimeStepsT1Device.rawPointer()); + setEngineIOStatus &= mContext->setTensorAddress(binding_names::kDenoisedTrajectory, mDenoisedDevice.rawPointer()); + if (mUsesMRope) + { + setEngineIOStatus &= mContext->setTensorAddress(binding_names::kRopeCosSin, mRopeCosSinDevice.rawPointer()); + } + setEngineIOStatus &= mContext->setTensorAddress(binding_names::kAttentionPosId, mPositionIdsDevice.rawPointer()); + + for (int32_t i = 0; i < mConfig.numDecoderLayers; ++i) + { + auto [kCacheBlock, vCacheBlock] = getSeparateKVCacheForDecoderLayer(stream, kvcache, i, activeBatchSize); + setEngineIOStatus + &= mContext->setTensorAddress(binding_names::formatKCacheName(i, true).c_str(), kCacheBlock.rawPointer()); + setEngineIOStatus + &= mContext->setTensorAddress(binding_names::formatVCacheName(i, true).c_str(), vCacheBlock.rawPointer()); + setEngineIOStatus + &= mContext->setTensorAddress(binding_names::formatKCacheName(i, false).c_str(), kCacheBlock.rawPointer()); + setEngineIOStatus + &= mContext->setTensorAddress(binding_names::formatVCacheName(i, false).c_str(), vCacheBlock.rawPointer()); + } + + if (!setEngineIOStatus) + { + LOG_ERROR("ActionRunner: failed to bind prefix-KV action engine tensors"); + return result; + } + + for (int32_t i = 0; i < mConfig.numInferenceTimesteps; ++i) + { + static_cast(mTimeStepsT0Host.rawPointer())[i] + = static_cast(i) / static_cast(mConfig.numInferenceTimesteps); + static_cast(mTimeStepsT1Host.rawPointer())[i] + = static_cast(i + 1) / static_cast(mConfig.numInferenceTimesteps); + } + + int32_t const* lengthsHost = getActualKVLengths(stream, activeBatchSize); + + int32_t* positionIdsPtr = static_cast(mPositionIdsHost.rawPointer()); + for (int32_t b = 0; b < activeBatchSize; ++b) + { + for (int32_t w = 0; w < mConfig.actionHorizon; ++w) + { + positionIdsPtr[b * mConfig.actionHorizon + w] = w; + } + } + + if (mUsesMRope) + { + int64_t* mropePosPtr = static_cast(mRopePositionIdsHost.rawPointer()); + for (int32_t b = 0; b < activeBatchSize; ++b) + { + int64_t const basePos = vlmOutputsRopeDeltas[b] + static_cast(lengthsHost[b]); + for (int32_t dim = 0; dim < 3; ++dim) + { + for (int32_t w = 0; w < mConfig.actionHorizon; ++w) + { + mropePosPtr[b * mConfig.actionHorizon * 3 + dim * mConfig.actionHorizon + w] + = static_cast(w) + basePos; + } + } + } + + CUDA_CHECK(cudaMemcpyAsync(mRopePositionIdsDevice.rawPointer(), mRopePositionIdsHost.rawPointer(), + static_cast(activeBatchSize) * 3ULL * static_cast(mConfig.actionHorizon) * sizeof(int64_t), + cudaMemcpyHostToDevice, stream)); + + kernel::initializeMRopeCosSin(static_cast(mRopeCosSinDevice.rawPointer()), + static_cast(mRopePositionIdsDevice.rawPointer()), mConfig.ropeTheta, + static_cast(mRopeHeadDim), mConfig.actionHorizon, activeBatchSize, true, mConfig.mropeSectionH, + mConfig.mropeSectionW, stream); + } + + CUDA_CHECK(cudaMemcpyAsync(mNoiseDevice.rawPointer(), mNoiseHost.rawPointer(), + static_cast(activeBatchSize) * static_cast(mConfig.actionHorizon) + * static_cast(mConfig.actionDim) * sizeof(float), + cudaMemcpyHostToDevice, stream)); + CUDA_CHECK(cudaMemcpyAsync(mPositionIdsDevice.rawPointer(), mPositionIdsHost.rawPointer(), + static_cast(activeBatchSize) * static_cast(mConfig.actionHorizon) * sizeof(int32_t), + cudaMemcpyHostToDevice, stream)); + + for (int32_t i = 0; i < mConfig.numInferenceTimesteps; ++i) + { + float* t0Ptr = static_cast(mTimeStepsT0Host.rawPointer()) + i; + float* t1Ptr = static_cast(mTimeStepsT1Host.rawPointer()) + i; + + CUDA_CHECK( + cudaMemcpyAsync(mTimeStepsT0Device.rawPointer(), t0Ptr, sizeof(float), cudaMemcpyHostToDevice, stream)); + CUDA_CHECK( + cudaMemcpyAsync(mTimeStepsT1Device.rawPointer(), t1Ptr, sizeof(float), cudaMemcpyHostToDevice, stream)); + + if (!mContext->enqueueV3(stream)) + { + LOG_ERROR("ActionRunner: enqueueV3 failed at denoise step %d", i); + return result; + } + + CUDA_CHECK(cudaMemcpyAsync(mNoiseDevice.rawPointer(), mDenoisedDevice.rawPointer(), + static_cast(activeBatchSize) * static_cast(mConfig.actionHorizon) + * static_cast(mConfig.actionDim) * sizeof(float), + cudaMemcpyDeviceToDevice, stream)); + } + + CUDA_CHECK(cudaMemcpyAsync(mDenoisedHost.rawPointer(), mDenoisedDevice.rawPointer(), + static_cast(activeBatchSize) * static_cast(mConfig.actionHorizon) + * static_cast(mConfig.actionDim) * sizeof(float), + cudaMemcpyDeviceToHost, stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + + result.resize(activeBatchSize); + float* output = static_cast(mDenoisedHost.rawPointer()); + for (int32_t b = 0; b < activeBatchSize; ++b) + { + result[b].reserve(mConfig.actionHorizon); + float const* row = output + b * (mConfig.actionHorizon * mConfig.actionDim); + for (int32_t w = 0; w < mConfig.actionHorizon; ++w) + { + result[b].emplace_back(row[w * mConfig.actionDim], row[w * mConfig.actionDim + 1]); + } + } + return result; +} + +std::size_t ActionRunner::velocityInputIndex(std::string const& name) const +{ + for (std::size_t i = 0; i < mInputNames.size(); ++i) + { + if (mInputNames[i] == name) + { + return i; + } + } + throw std::runtime_error("ActionRunner: unknown input tensor: " + name); +} + +bool ActionRunner::hasInputTensor(std::string const& name) const noexcept +{ + for (auto const& inputName : mInputNames) + { + if (inputName == name) + { + return true; + } + } + return false; +} + +rt::Tensor& ActionRunner::getActions() +{ + return mInputTensors.at(velocityInputIndex(mNoiseInputName)); +} + +rt::Tensor const& ActionRunner::getActions() const +{ + return mInputTensors.at(velocityInputIndex(mNoiseInputName)); +} + +bool ActionRunner::copyInputFrom(std::string const& name, rt::Tensor const& src, cudaStream_t stream) +{ + if (!hasInputTensor(name)) + { + LOG_ERROR("ActionRunner: unknown input tensor %s", name.c_str()); + return false; + } + auto& dst = mInputTensors.at(velocityInputIndex(name)); + return copyTensorToInput(dst, src, stream); +} + +bool ActionRunner::wireLanguageOutputs(std::vector const& languageOutputs, cudaStream_t stream) +{ + if (!mLmToActionSlots.empty()) + { + for (auto const& [lmSlot, actionSlot] : mLmToActionSlots) + { + if (lmSlot < 0 || static_cast(lmSlot) >= languageOutputs.size() + || languageOutputs[static_cast(lmSlot)] == nullptr) + { + LOG_ERROR( + "ActionRunner: invalid language output slot %d (num outputs=%zu)", lmSlot, languageOutputs.size()); + return false; + } + auto const& actionName = mInputNames[static_cast(actionSlot)]; + if (!copyInputFrom(actionName, *languageOutputs[static_cast(lmSlot)], stream)) + { + return false; + } + } + return true; + } + + if (hasInputTensor(mContextEmbedsName) && languageOutputs.size() == 1 && languageOutputs.front() != nullptr) + { + return copyInputFrom(mContextEmbedsName, *languageOutputs.front(), stream); + } + + return true; +} + +bool ActionRunner::isRolloutManagedInput(std::string const& name) const noexcept +{ + if (name == mNoiseInputName || name == mTimestepName) + { + return true; + } + for (auto const& wiredName : mLmWiredInputNames) + { + if (wiredName == name) + { + return true; + } + } + return false; +} + +bool ActionRunner::preparePi05SuffixInputs(int32_t activeBatchSize, int32_t prefixValidLen, cudaStream_t stream) +{ + if (mContextHandoff != ActionContextHandoff::CONTEXT_TENSOR) + { + return true; + } + if (!hasInputTensor("position_ids") || !hasInputTensor("attention_mask")) + { + return true; + } + + int32_t const prefixSeqLen = mConfigJson.value("prefix_seq_len", prefixValidLen); + int32_t const suffixLen = mConfig.actionHorizon; + if (suffixLen <= 0 || prefixSeqLen <= 0) + { + LOG_ERROR("ActionRunner::preparePi05SuffixInputs requires positive prefix_seq_len and action horizon"); + return false; + } + + int32_t const totalKeys = prefixSeqLen + suffixLen; + constexpr float kOpenPiAttentionMaskValue = -2.3819763e38F; + + auto& positionIdsTensor = mInputTensors.at(velocityInputIndex("position_ids")); + auto& attentionMaskTensor = mInputTensors.at(velocityInputIndex("attention_mask")); + + rt::Tensor positionIdsHost(positionIdsTensor.getShape(), rt::DeviceType::kCPU, positionIdsTensor.getDataType(), + "ActionRunner::preparePi05SuffixInputs::positionIdsHost"); + rt::Tensor attentionMaskHost(attentionMaskTensor.getShape(), rt::DeviceType::kCPU, + attentionMaskTensor.getDataType(), "ActionRunner::preparePi05SuffixInputs::attentionMaskHost"); + + std::vector suffixAttAr(static_cast(suffixLen), 0); + if (!suffixAttAr.empty()) + { + suffixAttAr[0] = 1; + } + std::vector suffixAttCumsum(static_cast(suffixLen), 0); + int64_t attCumsumRunning{0}; + for (int32_t idx = 0; idx < suffixLen; ++idx) + { + attCumsumRunning += suffixAttAr[static_cast(idx)]; + suffixAttCumsum[static_cast(idx)] = attCumsumRunning; + } + + int64_t* positionIdsData = positionIdsHost.dataPointer(); + float* attentionMaskData = attentionMaskHost.dataPointer(); + for (int32_t batchIdx = 0; batchIdx < activeBatchSize; ++batchIdx) + { + for (int32_t queryIdx = 0; queryIdx < suffixLen; ++queryIdx) + { + positionIdsData[(static_cast(batchIdx) * static_cast(suffixLen)) + + static_cast(queryIdx)] = static_cast(prefixValidLen + queryIdx); + + for (int32_t keyIdx = 0; keyIdx < totalKeys; ++keyIdx) + { + bool canAttend{false}; + if (keyIdx < prefixSeqLen) + { + canAttend = keyIdx < prefixValidLen; + } + else + { + int32_t const suffixKeyIdx = keyIdx - prefixSeqLen; + canAttend = suffixAttCumsum[static_cast(queryIdx)] + >= suffixAttCumsum[static_cast(suffixKeyIdx)]; + } + + std::size_t const maskOffset = (static_cast(batchIdx) * static_cast(suffixLen) + * static_cast(totalKeys)) + + (static_cast(queryIdx) * static_cast(totalKeys)) + + static_cast(keyIdx); + attentionMaskData[maskOffset] = canAttend ? 0.0F : kOpenPiAttentionMaskValue; + } + } + } + + CUDA_CHECK(cudaMemcpyAsync(positionIdsTensor.rawPointer(), positionIdsHost.rawPointer(), + tensorBytes(positionIdsTensor), cudaMemcpyHostToDevice, stream)); + CUDA_CHECK(cudaMemcpyAsync(attentionMaskTensor.rawPointer(), attentionMaskHost.rawPointer(), + tensorBytes(attentionMaskTensor), cudaMemcpyHostToDevice, stream)); + return true; +} + +bool ActionRunner::wireStaticInputs(LLMGenerationRequest const& request, cudaStream_t stream) +{ + if (mContextHandoff != ActionContextHandoff::CONTEXT_TENSOR) + { + return true; + } + + int32_t const numUniqueRequests = static_cast(request.requests.size()); + int32_t const actionBatchSize + = (request.actionBatchSize > 0) ? request.actionBatchSize : std::max(mActiveActionBatchSize, 1); + + auto resolveEmbodimentId = [&](LLMGenerationRequest::Request const& req) -> int64_t { + if (req.embodimentId.has_value()) + { + return req.embodimentId.value(); + } + if (request.embodimentId.has_value()) + { + return request.embodimentId.value(); + } + if (mConfigJson.contains("embodiment_id") && mConfigJson.at("embodiment_id").is_array() + && !mConfigJson.at("embodiment_id").empty()) + { + return mConfigJson.at("embodiment_id").at(0).get(); + } + return 31; + }; + + for (auto const& name : mInputNames) + { + if (isRolloutManagedInput(name)) + { + continue; + } + + auto& tensor = mInputTensors.at(velocityInputIndex(name)); + auto const elements = static_cast(tensor.getShape().volume()); + std::vector hostValues(elements, 0.0F); + + if (name == "state") + { + for (int32_t batchIdx = 0; batchIdx < actionBatchSize; ++batchIdx) + { + int32_t const reqIdx = numUniqueRequests > 0 ? batchIdx % numUniqueRequests : 0; + auto const& req = request.requests[static_cast(reqIdx)]; + if (!req.robotState.empty()) + { + if (req.robotState.size() != elements) + { + LOG_ERROR("ActionRunner: robot state size mismatch for batch %d (expected %zu, got %zu)", + batchIdx, elements, req.robotState.size()); + return false; + } + hostValues = req.robotState; + } + break; + } + } + else if (name == "embodiment_id") + { + int64_t const embodimentId = resolveEmbodimentId( + request.requests.empty() ? LLMGenerationRequest::Request{} : request.requests.front()); + hostValues.assign(elements, static_cast(embodimentId)); + } + + rt::Tensor hostTensor(tensor.getShape(), rt::DeviceType::kCPU, tensor.getDataType(), name + "_host"); + if (!fillHostTensorFromFloats(hostTensor, hostValues)) + { + return false; + } + CUDA_CHECK(cudaMemcpyAsync( + tensor.rawPointer(), hostTensor.rawPointer(), tensorBytes(tensor), cudaMemcpyHostToDevice, stream)); + } + + return true; +} + +bool ActionRunner::copyActionsToHost(std::vector>& actionsPerBatch, cudaStream_t stream) const +{ + if (mContextHandoff != ActionContextHandoff::CONTEXT_TENSOR) + { + return false; + } + + rt::Tensor const& actions = getActions(); + int32_t const batchSize = mActiveActionBatchSize > 0 ? mActiveActionBatchSize : mMaxActionBatchSize; + int64_t const perBatchElements = (actions.getShape().getNumDims() > 0) + ? actions.getShape().volume() / std::max(actions.getShape()[0], int64_t{1}) + : actions.getShape().volume(); + + rt::Tensor actionsHost( + actions.getShape(), rt::DeviceType::kCPU, actions.getDataType(), "ActionRunner::actionsHost"); + CUDA_CHECK(cudaMemcpyAsync( + actionsHost.rawPointer(), actions.rawPointer(), tensorBytes(actions), cudaMemcpyDeviceToHost, stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + + actionsPerBatch.clear(); + actionsPerBatch.resize(static_cast(batchSize)); + for (int32_t batchIdx = 0; batchIdx < batchSize; ++batchIdx) + { + actionsPerBatch[static_cast(batchIdx)].resize(static_cast(perBatchElements)); + auto& dst = actionsPerBatch[static_cast(batchIdx)]; + if (actions.getDataType() == nvinfer1::DataType::kFLOAT) + { + auto const* src = actionsHost.dataPointer() + + static_cast(batchIdx) * static_cast(perBatchElements); + std::copy(src, src + static_cast(perBatchElements), dst.begin()); + } + else if (actions.getDataType() == nvinfer1::DataType::kHALF) + { + auto const* src = actionsHost.dataPointer() + + static_cast(batchIdx) * static_cast(perBatchElements); + for (int64_t i = 0; i < perBatchElements; ++i) + { + dst[static_cast(i)] = __half2float(src[static_cast(i)]); + } + } + else + { + LOG_ERROR("ActionRunner: unsupported action dtype for export"); + return false; + } + } + + return true; +} + +bool ActionRunner::bindVelocityTensors() noexcept +{ + for (std::size_t i = 0; i < mInputNames.size(); ++i) + { + if (!bindTensor(mEngine.get(), mContext.get(), engineBindingName(i), mInputTensors[i])) + { + return false; + } + } + return bindTensor(mEngine.get(), mContext.get(), mPredVelocityName, mPredVelocity); +} + +void ActionRunner::buildEngineBindingNames() +{ + std::vector engineInputNames; + for (int32_t i = 0; i < mEngine->getNbIOTensors(); ++i) + { + char const* name = mEngine->getIOTensorName(i); + if (name != nullptr && mEngine->getTensorIOMode(name) == nvinfer1::TensorIOMode::kINPUT) + { + engineInputNames.emplace_back(name); + } + } + + if (engineInputNames.size() < mInputNames.size()) + { + throw std::runtime_error("ActionRunner: engine has fewer inputs than config input_names"); + } + + mEngineBindingNames.clear(); + mEngineBindingNames.reserve(mInputNames.size()); + for (std::size_t i = 0; i < mInputNames.size(); ++i) + { + auto const& logicalName = mInputNames[i]; + if (engineHasTensor(logicalName)) + { + mEngineBindingNames.push_back(logicalName); + } + else + { + mEngineBindingNames.push_back(engineInputNames[i]); + LOG_WARNING("ActionRunner: logical input '%s' mapped to engine binding '%s'", logicalName.c_str(), + engineInputNames[i].c_str()); + } + } +} + +std::string const& ActionRunner::engineBindingName(std::size_t inputIndex) const +{ + return mEngineBindingNames.at(inputIndex); +} + +bool ActionRunner::setTimestepForStep(int32_t step, cudaStream_t stream) +{ + auto const elements = mTimestepHost.getShape().volume(); + + if (mTimestepSchedule == "continuous_flow") + { + if (mTimestepHost.getDataType() != nvinfer1::DataType::kFLOAT) + { + LOG_ERROR("ActionRunner: continuous_flow schedule requires float32 timestep"); + return false; + } + + auto const dt = -1.0F / static_cast(mConfig.numInferenceTimesteps); + auto const timestepValue = 1.0F + static_cast(step) * dt; + auto* data = mTimestepHost.dataPointer(); + for (int64_t i = 0; i < elements; ++i) + { + data[i] = timestepValue; + } + } + else + { + check::check(mNumTimestepBuckets > 0, "ActionRunner num_timestep_buckets must be positive"); + auto const bucket = static_cast(std::floor(static_cast(step) + / static_cast(mConfig.numInferenceTimesteps) * static_cast(mNumTimestepBuckets))); + + if (mTimestepHost.getDataType() == nvinfer1::DataType::kINT64) + { + auto* data = mTimestepHost.dataPointer(); + for (int64_t i = 0; i < elements; ++i) + { + data[i] = bucket; + } + } + else if (mTimestepHost.getDataType() == nvinfer1::DataType::kINT32) + { + auto* data = mTimestepHost.dataPointer(); + for (int64_t i = 0; i < elements; ++i) + { + data[i] = static_cast(bucket); + } + } + else + { + LOG_ERROR("ActionRunner: discrete_buckets schedule requires int32 or int64 timestep"); + return false; + } + } + + auto& timestepTensor = mInputTensors.at(velocityInputIndex(mTimestepName)); + CUDA_CHECK(cudaMemcpyAsync(timestepTensor.rawPointer(), mTimestepHost.rawPointer(), tensorBytes(timestepTensor), + cudaMemcpyHostToDevice, stream)); + return true; +} + +bool ActionRunner::updateActionsOnHost(cudaStream_t stream, float stepSize) +{ + CUDA_CHECK(cudaMemcpyAsync(mPredVelocityHost.rawPointer(), mPredVelocity.rawPointer(), tensorBytes(mPredVelocity), + cudaMemcpyDeviceToHost, stream)); + CUDA_CHECK(cudaMemcpyAsync( + mNoiseHost.rawPointer(), getActions().rawPointer(), tensorBytes(getActions()), cudaMemcpyDeviceToHost, stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + + updateActionTensor(mNoiseHost, mPredVelocityHost, stepSize); + + CUDA_CHECK(cudaMemcpyAsync( + getActions().rawPointer(), mNoiseHost.rawPointer(), tensorBytes(getActions()), cudaMemcpyHostToDevice, stream)); + return true; +} + +bool ActionRunner::sampleActions(cudaStream_t stream) +{ + if (mContextHandoff != ActionContextHandoff::CONTEXT_TENSOR) + { + LOG_ERROR("ActionRunner::sampleActions requires CONTEXT_TENSOR handoff"); + return false; + } + + check::check(mConfig.numInferenceTimesteps > 0, "ActionRunner num_inference_timesteps must be positive"); + + int32_t const activeBatchSize = mActiveActionBatchSize > 0 ? mActiveActionBatchSize : mMaxActionBatchSize; + if (!reshapeActionTensorsForActiveBatch(activeBatchSize)) + { + return false; + } + + CUDA_CHECK(cudaMemcpyAsync( + getActions().rawPointer(), mNoiseHost.rawPointer(), tensorBytes(getActions()), cudaMemcpyHostToDevice, stream)); + + auto const stepSize = static_cast(mRolloutDtSign) / static_cast(mConfig.numInferenceTimesteps); + + for (int32_t step = 0; step < mConfig.numInferenceTimesteps; ++step) + { + if (!setTimestepForStep(step, stream) || !bindVelocityTensors() || !mContext->enqueueV3(stream) + || !updateActionsOnHost(stream, stepSize)) + { + return false; + } + } + + return true; +} + +} // namespace rt +} // namespace trt_edgellm diff --git a/cpp/action/actionRunner.h b/cpp/action/actionRunner.h new file mode 100644 index 00000000..67dd152c --- /dev/null +++ b/cpp/action/actionRunner.h @@ -0,0 +1,224 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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. + */ + +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "action/actionModelTypes.h" +#include "common/tensor.h" +#include "runtime/linearKVCache.h" +#include "runtime/llmRuntimeUtils.h" +#include "tokenizer/tokenizer.h" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace trt_edgellm +{ +namespace rt +{ + +//! How frozen VLM/LM context is passed into the action engine for a request. +enum class ActionContextHandoff +{ + PREFIX_KV, //!< Per-layer k_cache_i / v_cache_i from LinearKVCache (e.g. Alpamayo) + CONTEXT_TENSOR, //!< context_embs and related tensors wired from language prefill (e.g. GR00T) + UNKNOWN +}; + +//! How x_t is updated after each denoise engine forward. +enum class ActionRolloutMode +{ + FLOW_MATCHING, //!< Engine outputs next state; copy denoised -> noise + VELOCITY, //!< Engine outputs velocity; integrate x_t += dt * velocity on host + UNKNOWN +}; + +//! Fields loaded from action/config.json. Not every field applies to every model. +struct ActionRunnerConfig +{ + float ropeTheta{0.0F}; + int32_t mropeSectionH{0}; + int32_t mropeSectionW{0}; + int32_t numDecoderLayers{0}; + int32_t numTrajTokens{0}; + int32_t trajTokenStart{0}; + int32_t maxKVCacheCapacity{0}; + int32_t numInferenceTimesteps{0}; + int32_t actionHorizon{0}; + int32_t actionDim{0}; +}; + +//! Model-agnostic orchestration for a compiled action / diffusion step engine. +class ActionRunner +{ +public: + ActionRunner(std::string const& engineDir, cudaStream_t stream, LinearKVCache::CacheConfig const& kvCacheConfig); + + ~ActionRunner() noexcept = default; + + int64_t getRequiredContextMemorySize() const; + bool setContextMemory(rt::Tensor& sharedContextMemory); + bool resetExecutionContext(cudaStream_t stream); + + action::ActionModelType getModelType() const noexcept + { + return mModelType; + } + + ActionContextHandoff getContextHandoff() const noexcept + { + return mContextHandoff; + } + + ActionRolloutMode getRolloutMode() const noexcept + { + return mRolloutMode; + } + + void setNoiseSeed(int32_t seed) noexcept + { + mNoiseSeed = seed; + mRng.seed(static_cast(seed)); + } + + int32_t getMaxKVCacheCapacity() const noexcept + { + return mConfig.maxKVCacheCapacity; + } + + bool preprocess(LLMGenerationRequest const& request, std::vector>& batchedInputIds, + tokenizer::Tokenizer const* tokenizer); + + //! Prefix-KV / flow-matching path (Alpamayo-style). Requires populated kvcache and optional MRoPE deltas. + std::vector> sampleTrajectory(cudaStream_t stream, int32_t activeBatchSize, + LinearKVCache& kvcache, std::vector const& vlmOutputsRopeDeltas); + + //! Context-tensor / velocity path (GR00T-style). Wire context via copyInputFrom / wireLanguageOutputs first. + bool sampleActions(cudaStream_t stream); + + bool hasInputTensor(std::string const& name) const noexcept; + bool copyInputFrom(std::string const& name, rt::Tensor const& src, cudaStream_t stream); + bool wireLanguageOutputs(std::vector const& languageOutputs, cudaStream_t stream); + //! Bind non-rollout action inputs (state, embodiment_id, etc.) from the request. + bool wireStaticInputs(LLMGenerationRequest const& request, cudaStream_t stream); + //! Build PI0.5 suffix position_ids / attention_mask from the packed prefix length. + bool preparePi05SuffixInputs(int32_t activeBatchSize, int32_t prefixValidLen, cudaStream_t stream); + bool isRolloutManagedInput(std::string const& name) const noexcept; + rt::Tensor& getActions(); + rt::Tensor const& getActions() const; + bool copyActionsToHost(std::vector>& actionsPerBatch, cudaStream_t stream) const; + +private: + bool parseModelConfig(std::string const& configPath); + void allocateTensors(LinearKVCache::CacheConfig const& kvCacheConfig); + bool reshapeActionTensorsForActiveBatch(int32_t activeBatchSize); + void initializeNoiseTrajectory(int32_t randomSeed, int32_t activeBatchSize); + void setDynamicInputShapes(int32_t activeBatchSize); + int32_t const* getActualKVLengths(cudaStream_t stream, int32_t activeBatchSize); + std::pair getSeparateKVCacheForDecoderLayer( + cudaStream_t stream, LinearKVCache& kvcache, int32_t decoderLayerIdx, int32_t activeBatchSize); + + bool engineHasTensor(std::string const& name) const noexcept; + void allocatePrefixKVTensors(LinearKVCache::CacheConfig const& kvCacheConfig); + void allocateVelocityTensors(); + bool bindVelocityTensors() noexcept; + bool setTimestepForStep(int32_t step, cudaStream_t stream); + bool updateActionsOnHost(cudaStream_t stream, float stepSize); + void buildEngineBindingNames(); + std::size_t velocityInputIndex(std::string const& name) const; + std::string const& engineBindingName(std::size_t inputIndex) const; + + static constexpr int32_t kDefaultDenoiseSteps = 10; + + cudaStream_t mStream{nullptr}; + action::ActionModelType mModelType{action::ActionModelType::UNKNOWN}; + ActionContextHandoff mContextHandoff{ActionContextHandoff::UNKNOWN}; + ActionRolloutMode mRolloutMode{ActionRolloutMode::UNKNOWN}; + int32_t mNoiseSeed{5}; + ActionRunnerConfig mConfig{}; + nlohmann::json mConfigJson; + + std::unique_ptr mRuntime{nullptr}; + std::unique_ptr mEngine{nullptr}; + std::unique_ptr mContext{nullptr}; + void* mExecContextMemory{nullptr}; + int64_t mExecContextMemoryCapacity{0}; + + int32_t mMaxActionBatchSize{0}; + int32_t mActiveActionBatchSize{0}; + bool mUsesMRope{false}; + + // Shared rollout noise [B, horizon, dim] + rt::Tensor mNoiseHost; + rt::Tensor mNoiseDevice; + + // Prefix-KV / flow-matching tensors + rt::Tensor mDenoisedDevice; + rt::Tensor mDenoisedHost; + rt::Tensor mTimeStepsT0Device; + rt::Tensor mTimeStepsT1Device; + rt::Tensor mTimeStepsT0Host; + rt::Tensor mTimeStepsT1Host; + rt::Tensor mKvcacheActualLengthsHost; + int32_t* mKvcacheActualLengthsDevice{nullptr}; + rt::Tensor mKvcacheActualLengthsBroadcastDevice; + rt::Tensor mRopeCosSinDevice; + rt::Tensor mPositionIdsHost; + rt::Tensor mPositionIdsDevice; + rt::Tensor mRopePositionIdsHost; + rt::Tensor mRopePositionIdsDevice; + int32_t mNumKVHeads{0}; + int32_t mMaxSequenceLength{0}; + int32_t mKvHeadDim{0}; + int64_t mRopeHeadDim{0}; + std::vector mKCacheLayers; + std::vector mVCacheLayers; + + // Context-tensor / velocity tensors + std::vector mInputNames; + std::vector mEngineBindingNames; + std::vector mInputTensors; + rt::Tensor mPredVelocity; + rt::Tensor mPredVelocityHost; + rt::Tensor mTimestepHost; + std::string mNoiseInputName{"actions"}; + std::string mTimestepName{"timestep"}; + std::string mContextEmbedsName{"context_embs"}; + std::string mPredVelocityName{"velocity"}; + std::string mTimestepSchedule{"discrete_buckets"}; + std::vector> mLmToActionSlots; + std::vector mLmWiredInputNames; + int32_t mNumTimestepBuckets{1}; + int32_t mRolloutDtSign{1}; + std::mt19937 mRng{0}; +}; + +} // namespace rt +} // namespace trt_edgellm diff --git a/cpp/common/bindingNames.h b/cpp/common/bindingNames.h index d2c826ed..b60e2f46 100644 --- a/cpp/common/bindingNames.h +++ b/cpp/common/bindingNames.h @@ -79,6 +79,34 @@ inline constexpr char const* kLogits = "logits"; */ inline constexpr char const* kOutputHiddenStates = "hidden_states"; +/*! + * @brief GR00T context embeddings for VLA action head + * + * Shape: [batch_size, sequence_length, context_hidden_size] (FLOAT16) + */ +inline constexpr char const* kOutputContextEmbeds = "context_embs"; + +/*! + * @brief Language-model hidden states for split GR00T action-context engine + * + * Shape: [batch_size, sequence_length, hidden_size] (FLOAT16) + */ +inline constexpr char const* kOutputLmHiddenStates = "lm_hidden_states"; + +/*! + * @brief PI0.5 prefix key cache stacked across decoder layers + * + * Shape: [num_layers, batch_size, num_kv_heads, 1, prefix_seq_len, head_dim] (FLOAT16) + */ +inline constexpr char const* kOutputPrefixK = "prefix_k"; + +/*! + * @brief PI0.5 prefix value cache stacked across decoder layers + * + * Shape: [num_layers, batch_size, num_kv_heads, 1, prefix_seq_len, head_dim] (FLOAT16) + */ +inline constexpr char const* kOutputPrefixV = "prefix_v"; + /*! * @brief DFlash draft model input: concatenated target hidden states. * @@ -331,6 +359,22 @@ inline constexpr char const* kVisualInput = "input"; */ inline constexpr char const* kVisualOutput = "output"; +/*! + * @brief Visual input tensor for the model-agnostic fixed-shape VitRunner. + * + * VitRunner engines are exported with dedicated binding names distinct from the + * shared Qwen-VL / InternVL `kVisualInput`/`kVisualOutput`. + * Shape: [batch, height, width, channels] HWC (FLOAT16) + */ +inline constexpr char const* kVitInput = "pixel_values"; + +/*! + * @brief Visual output tensor for the model-agnostic fixed-shape VitRunner. + * + * Shape: [num_image_tokens, hidden_size] (FLOAT16) + */ +inline constexpr char const* kVitOutput = "visual_embeds"; + /*! * @brief Rotary positional embeddings for visual inputs * diff --git a/cpp/kernels/kvCacheUtilKernels/kvCacheUtilsKernels.cu b/cpp/kernels/kvCacheUtilKernels/kvCacheUtilsKernels.cu index 6f5f365c..1eaaf487 100644 --- a/cpp/kernels/kvCacheUtilKernels/kvCacheUtilsKernels.cu +++ b/cpp/kernels/kvCacheUtilKernels/kvCacheUtilsKernels.cu @@ -380,5 +380,115 @@ void instantiateKVCacheBatched(KVLayerInfo const* dstLayerInfos, KVLayerInfo con CUDA_CHECK(cudaGetLastError()); } +void instantiateKVCacheFromTensor( + rt::Tensor& dstKVCacheBuffer, rt::Tensor const& srcKVCacheTensor, int32_t batchIdx, cudaStream_t stream) +{ + // srcKVCacheTensor shape: [numDecoderLayers, 2, numKVHeads, sequenceLength, headDim] + int32_t const numDecoderLayers = srcKVCacheTensor.getShape()[0]; + int32_t const numKVHeads = srcKVCacheTensor.getShape()[2]; + int32_t const sequenceLength = srcKVCacheTensor.getShape()[3]; + int32_t const headDim = srcKVCacheTensor.getShape()[4]; + + // dstKVCacheBuffer shape: [numDecoderLayers, maxBatchSize, 2, numKVHeads, maxSequenceLength, headDim] + int32_t const kvCacheMaxBatch = dstKVCacheBuffer.getShape()[1]; + int32_t const kvCacheMaxSequenceLength = dstKVCacheBuffer.getShape()[4]; + + if (batchIdx >= kvCacheMaxBatch) + { + throw std::runtime_error( + "instantiateKVCacheFromTensor(): batchIdx is out of range for the KVCache buffer. MaxSupportedBatch = " + + std::to_string(kvCacheMaxBatch) + ", batchIdx = " + std::to_string(batchIdx)); + } + if (sequenceLength > kvCacheMaxSequenceLength) + { + throw std::runtime_error( + "instantiateKVCacheFromTensor(): sequenceLength is out of range for the KVCache buffer. " + "MaxSupportedSequenceLength = " + + std::to_string(kvCacheMaxSequenceLength) + ", sequenceLength = " + std::to_string(sequenceLength)); + } + if (dstKVCacheBuffer.getDataType() != srcKVCacheTensor.getDataType() + && dstKVCacheBuffer.getDataType() != nvinfer1::DataType::kHALF) + { + throw std::runtime_error( + "instantiateKVCacheFromTensor(): KVCacheBuffer and preComputedKVCache shall both be half type now."); + } + + dim3 gridDim(numDecoderLayers * 2 * numKVHeads); + dim3 blockDim(32, 4); + half* srcKVCacheTensorPtr = const_cast(srcKVCacheTensor.dataPointer()); + switch (headDim) + { + case 64: + instantiateKVCacheKernel<<>>(dstKVCacheBuffer.dataPointer(), + srcKVCacheTensorPtr, kvCacheMaxBatch, kvCacheMaxSequenceLength, batchIdx, numDecoderLayers, numKVHeads, + sequenceLength, headDim); + break; + case 128: + instantiateKVCacheKernel<<>>( + dstKVCacheBuffer.dataPointer(), srcKVCacheTensorPtr, kvCacheMaxBatch, kvCacheMaxSequenceLength, + batchIdx, numDecoderLayers, numKVHeads, sequenceLength, headDim); + break; + default: + throw std::runtime_error( + "instantiateKVCacheFromTensor(): Only headDim = 64 or 128 are supported by the kernel, current headDim = " + + std::to_string(headDim)); + } + CUDA_CHECK(cudaGetLastError()); +} + +void saveKVCacheIntoTensor( + rt::Tensor& dstKVCacheTensor, rt::Tensor const& srcKVCacheBuffer, int32_t batchIdx, cudaStream_t stream) +{ + int32_t const numDecoderLayers = dstKVCacheTensor.getShape()[0]; + int32_t const numKVHeads = dstKVCacheTensor.getShape()[2]; + int32_t const sequenceLength = dstKVCacheTensor.getShape()[3]; + int32_t const headDim = dstKVCacheTensor.getShape()[4]; + + int32_t const kvCacheMaxBatch = srcKVCacheBuffer.getShape()[1]; + int32_t const kvCacheMaxSequenceLength = srcKVCacheBuffer.getShape()[4]; + + if (batchIdx >= kvCacheMaxBatch) + { + throw std::runtime_error( + "saveKVCacheIntoTensor(): batchIdx is out of range for the KVCache buffer. MaxSupportedBatch = " + + std::to_string(kvCacheMaxBatch) + ", batchIdx = " + std::to_string(batchIdx)); + } + if (sequenceLength > kvCacheMaxSequenceLength) + { + throw std::runtime_error( + "saveKVCacheIntoTensor(): sequenceLength is out of range for the KVCache buffer. " + "MaxSupportedSequenceLength = " + + std::to_string(kvCacheMaxSequenceLength) + ", sequenceLength = " + std::to_string(sequenceLength)); + } + if (dstKVCacheTensor.getDataType() != srcKVCacheBuffer.getDataType() + && dstKVCacheTensor.getDataType() != nvinfer1::DataType::kHALF) + { + throw std::runtime_error( + "saveKVCacheIntoTensor(): KVCacheBuffer and preComputedKVCache shall both be half type now."); + } + + dim3 gridDim(numDecoderLayers * 2 * numKVHeads); + dim3 blockDim(32, 4); + half* srcKVCacheBufferPtr = const_cast(srcKVCacheBuffer.dataPointer()); + switch (headDim) + { + case 64: + instantiateKVCacheKernel<<>>(srcKVCacheBufferPtr, + dstKVCacheTensor.dataPointer(), kvCacheMaxBatch, kvCacheMaxSequenceLength, batchIdx, numDecoderLayers, + numKVHeads, sequenceLength, headDim); + break; + case 128: + instantiateKVCacheKernel<<>>(srcKVCacheBufferPtr, + dstKVCacheTensor.dataPointer(), kvCacheMaxBatch, kvCacheMaxSequenceLength, batchIdx, numDecoderLayers, + numKVHeads, sequenceLength, headDim); + break; + default: + throw std::runtime_error( + "saveKVCacheIntoTensor(): Only headDim = 64 or 128 are supported by the kernel, current headDim = " + + std::to_string(headDim)); + } + CUDA_CHECK(cudaGetLastError()); +} + } // namespace kernel } // namespace trt_edgellm \ No newline at end of file diff --git a/cpp/kernels/kvCacheUtilKernels/kvCacheUtilsKernels.h b/cpp/kernels/kvCacheUtilKernels/kvCacheUtilsKernels.h index bc6a633a..cf051182 100644 --- a/cpp/kernels/kvCacheUtilKernels/kvCacheUtilsKernels.h +++ b/cpp/kernels/kvCacheUtilKernels/kvCacheUtilsKernels.h @@ -65,6 +65,25 @@ void instantiateKVCacheLayerFromTensor( void saveKVCacheLayerIntoTensor( rt::Tensor& dstKVCacheTensor, rt::Tensor const& srcKVCacheLayer, int32_t batchIdx, cudaStream_t stream); +//! \brief Monolithic variant: instantiate the whole KV cache buffer from a saved tensor. +//! +//! Used by the VLA runtime (LLMEngineRunner) whose KV cache is a single contiguous buffer. +//! \param[in,out] dstKVCacheBuffer [numDecoderLayers, maxBatchSize, 2, numKVHeads, maxSequenceLength, headDim] +//! \param[in] srcKVCacheTensor [numDecoderLayers, 2, numKVHeads, sequenceLength, headDim] +//! \param[in] batchIdx Target batch index in the destination buffer +//! \param[in] stream CUDA stream +void instantiateKVCacheFromTensor( + rt::Tensor& dstKVCacheBuffer, rt::Tensor const& srcKVCacheTensor, int32_t batchIdx, cudaStream_t stream); + +//! \brief Monolithic variant: save the whole KV cache buffer into a tensor. +//! +//! \param[out] dstKVCacheTensor [numDecoderLayers, 2, numKVHeads, sequenceLength, headDim] +//! \param[in] srcKVCacheBuffer [numDecoderLayers, maxBatchSize, 2, numKVHeads, maxSequenceLength, headDim] +//! \param[in] batchIdx Source batch index in the buffer +//! \param[in] stream CUDA stream +void saveKVCacheIntoTensor( + rt::Tensor& dstKVCacheTensor, rt::Tensor const& srcKVCacheBuffer, int32_t batchIdx, cudaStream_t stream); + /// @brief Batched save: copy multiple layers' KV cache into per-layer tensors in a single launch. /// All layers must share the same headDim. dstLayerInfos[i].data points to a [2, numKVHeads_i, seqLen, headDim] tensor. /// @param srcLayerInfos [numLayers] GPU array — source cache buffers diff --git a/cpp/multimodal/modelTypes.h b/cpp/multimodal/modelTypes.h index 2fe77fe8..466ca411 100644 --- a/cpp/multimodal/modelTypes.h +++ b/cpp/multimodal/modelTypes.h @@ -39,6 +39,7 @@ enum class ModelType GEMMA4_VISION, //!< Gemma4 vision encoder NEMOTRON_OMNI_VISION_ENCODER, //!< Nemotron-Omni vision encoder NEMOTRON_OMNI_AUDIO_ENCODER, //!< Nemotron-Omni audio encoder + VIT, //!< Model-agnostic fixed-shape ViT (VisualFixedInput export) UNKNOWN //!< Unknown or unsupported model type }; @@ -72,6 +73,8 @@ inline ModelType stringToModelType(std::string const& modelTypeStr) return ModelType::NEMOTRON_OMNI_VISION_ENCODER; if (modelTypeStr == "parakeet") return ModelType::NEMOTRON_OMNI_AUDIO_ENCODER; + if (modelTypeStr == "vit") + return ModelType::VIT; return ModelType::UNKNOWN; } diff --git a/cpp/multimodal/multimodalRunner.cpp b/cpp/multimodal/multimodalRunner.cpp index 4192ee14..ef0bac3d 100644 --- a/cpp/multimodal/multimodalRunner.cpp +++ b/cpp/multimodal/multimodalRunner.cpp @@ -25,6 +25,7 @@ #include "multimodal/nemotronOmniViTRunner.h" #include "multimodal/phi4mmViTRunner.h" #include "multimodal/qwenViTRunner.h" +#include "multimodal/vitRunner.h" #include "profiling/layerProfiler.h" #include "profiling/metrics.h" #include "profiling/timer.h" @@ -155,6 +156,10 @@ std::unique_ptr MultimodalRunner::create(std::string const& mu { multimodalRunner = std::make_unique(multimodalEngineDir, stream); } + else if (modelType == multimodal::ModelType::VIT) + { + multimodalRunner = std::make_unique(multimodalEngineDir, stream); + } else { throw std::runtime_error("Unsupported model type: " + modelTypeStr); diff --git a/cpp/multimodal/vitRunner.cpp b/cpp/multimodal/vitRunner.cpp new file mode 100644 index 00000000..9c95f441 --- /dev/null +++ b/cpp/multimodal/vitRunner.cpp @@ -0,0 +1,317 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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. + */ + +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "multimodal/vitRunner.h" + +#include "common/bindingNames.h" +#include "common/checkMacros.h" +#include "common/logger.h" +#include "kernels/preprocessKernels/imageUtilKernels.h" +#include "multimodal/modelTypes.h" +#include "profiling/metrics.h" +#include "profiling/timer.h" + +#include +#include +#include +#include +#include + +using Json = nlohmann::json; + +namespace trt_edgellm +{ +namespace rt +{ + +VitRunner::VitRunner(std::string const& engineDir, cudaStream_t stream) + : MultimodalRunner(engineDir, stream) +{ + std::string const configPath = engineDir + "/config.json"; + if (!validateAndFillConfig(configPath)) + { + LOG_ERROR("VitRunner: failed to validate and fill config"); + throw std::runtime_error("VitRunner: failed to validate and fill config"); + } + if (!allocateBuffer(stream)) + { + LOG_ERROR("VitRunner: failed to allocate buffer"); + throw std::runtime_error("VitRunner: failed to allocate buffer"); + } +} + +bool VitRunner::validateAndFillConfig(std::string const& configPath) +{ + Json jsonConfig; + + std::ifstream configFileStream(configPath); + if (!configFileStream.is_open()) + { + LOG_ERROR("VitRunner: failed to open config file: %s", configPath.c_str()); + return false; + } + + try + { + jsonConfig = Json::parse(configFileStream); + configFileStream.close(); + } + catch (Json::parse_error const& e) + { + LOG_ERROR("VitRunner: failed to parse config file: %s", e.what()); + return false; + } + + std::string const modelTypeStr = jsonConfig["model_type"].get(); + mModelType = multimodal::stringToModelType(modelTypeStr); + if (mModelType != multimodal::ModelType::VIT) + { + LOG_ERROR("VitRunner: invalid model type: %s", modelTypeStr.c_str()); + return false; + } + + mConfig.vocabSize = jsonConfig["vocab_size"].get(); + mConfig.imageTokenId = jsonConfig["image_token_id"].get(); + + auto const builderConfig = jsonConfig["builder_config"]; + mConfig.seqLenPerImage = builderConfig["seq_len"].get(); + if (builderConfig.contains("image_mean")) + { + for (size_t i = 0; i < mConfig.imageMean.size(); ++i) + { + mConfig.imageMean[i] = builderConfig["image_mean"].at(i).get(); + } + } + if (builderConfig.contains("image_std")) + { + for (size_t i = 0; i < mConfig.imageStd.size(); ++i) + { + mConfig.imageStd[i] = builderConfig["image_std"].at(i).get(); + } + } + + // Get config from engine shapes ([batch, H, W, C] HWC input) + nvinfer1::Dims const inputShapeMax + = mVisualEngine->getProfileShape(binding_names::kVitInput, 0, nvinfer1::OptProfileSelector::kMAX); + mConfig.batchSize = inputShapeMax.d[0]; + mConfig.imageHeight = inputShapeMax.d[1]; + mConfig.imageWidth = inputShapeMax.d[2]; + mConfig.numChannels = inputShapeMax.d[3]; + + nvinfer1::Dims const outputShape = mVisualEngine->getTensorShape(binding_names::kVitOutput); + mConfig.numImageTokens = outputShape.d[0]; + mConfig.outHiddenSize = outputShape.d[1]; + + return true; +} + +bool VitRunner::allocateBuffer(cudaStream_t stream) +{ + bool setTensorAddressStatus{true}; + + mVitInput = rt::Tensor({mConfig.batchSize, mConfig.imageHeight, mConfig.imageWidth, mConfig.numChannels}, + rt::DeviceType::kGPU, nvinfer1::DataType::kHALF, "VitRunner::mVitInput"); + mOutputEmbedding = rt::Tensor({mConfig.numImageTokens, mConfig.outHiddenSize}, rt::DeviceType::kGPU, + nvinfer1::DataType::kHALF, "VitRunner::mOutputEmbedding"); + + setTensorAddressStatus &= mVisualContext->setTensorAddress(binding_names::kVitInput, mVitInput.rawPointer()); + setTensorAddressStatus + &= mVisualContext->setTensorAddress(binding_names::kVitOutput, mOutputEmbedding.rawPointer()); + if (!setTensorAddressStatus) + { + LOG_ERROR("VitRunner: failed to set tensor addresses on the engine"); + return false; + } + + int64_t const channels = static_cast(mConfig.numChannels); + mImageMean = rt::Tensor({channels}, rt::DeviceType::kGPU, nvinfer1::DataType::kFLOAT, "VitRunner::mImageMean"); + mImageStd = rt::Tensor({channels}, rt::DeviceType::kGPU, nvinfer1::DataType::kFLOAT, "VitRunner::mImageStd"); + CUDA_CHECK(cudaMemcpyAsync( + mImageMean.rawPointer(), mConfig.imageMean.data(), channels * sizeof(float), cudaMemcpyHostToDevice, stream)); + CUDA_CHECK(cudaMemcpyAsync( + mImageStd.rawPointer(), mConfig.imageStd.data(), channels * sizeof(float), cudaMemcpyHostToDevice, stream)); + + int64_t const maxImagePixels = mConfig.imageHeight * mConfig.imageWidth * mConfig.numChannels; + mImageDevice + = rt::Tensor({maxImagePixels}, rt::DeviceType::kGPU, nvinfer1::DataType::kUINT8, "VitRunner::mImageDevice"); + mNormalizedImageDevice = rt::Tensor( + {maxImagePixels}, rt::DeviceType::kGPU, nvinfer1::DataType::kHALF, "VitRunner::mNormalizedImageDevice"); + + rt::Tensor resizeBuffer( + {1, maxImagePixels, channels}, rt::DeviceType::kCPU, nvinfer1::DataType::kUINT8, "VitRunner::resizeBuffer"); + mResizedImageHost = rt::imageUtils::ImageData(std::move(resizeBuffer)); + + return true; +} + +void VitRunner::formatImage(rt::imageUtils::ImageData const& image, int64_t batchIndex, cudaStream_t stream) +{ + int64_t height = image.height; + int64_t width = image.width; + int64_t channels = image.channels; + if (channels != mConfig.numChannels) + { + throw std::runtime_error("VitRunner::formatImage(): image channels mismatch, got " + std::to_string(channels) + + ", expected " + std::to_string(mConfig.numChannels)); + } + + rt::imageUtils::ImageData const* imageToUse = ℑ + if (height != mConfig.imageHeight || width != mConfig.imageWidth) + { + rt::imageUtils::resizeImage(image, mResizedImageHost, mConfig.imageWidth, mConfig.imageHeight, + rt::imageUtils::InterpolationMode::kLINEAR); + imageToUse = &mResizedImageHost; + height = imageToUse->height; + width = imageToUse->width; + } + + check::check( + mImageDevice.reshape({1, height, width, channels}), "VitRunner::formatImage(): mImageDevice reshape failed"); + check::check(mNormalizedImageDevice.reshape({1, height, width, channels}), + "VitRunner::formatImage(): mNormalizedImageDevice reshape failed"); + + CUDA_CHECK(cudaMemcpyAsync( + mImageDevice.rawPointer(), imageToUse->data(), height * width * channels, cudaMemcpyHostToDevice, stream)); + kernel::normalizeImage(mImageDevice, mImageMean, mImageStd, mNormalizedImageDevice, stream); + + int64_t const imageBytes = height * width * channels * static_cast(sizeof(half)); + int64_t const batchOffset = batchIndex * imageBytes; + CUDA_CHECK(cudaMemcpyAsync(static_cast(mVitInput.rawPointer()) + batchOffset, + mNormalizedImageDevice.rawPointer(), static_cast(imageBytes), cudaMemcpyDeviceToDevice, stream)); +} + +void VitRunner::imagePreprocess(rt::LLMGenerationRequest const& request, cudaStream_t stream) +{ + int64_t batchIndex = 0; + for (auto const& req : request.requests) + { + for (auto const& image : req.imageBuffers) + { + if (batchIndex >= mConfig.batchSize) + { + throw std::runtime_error("VitRunner::imagePreprocess(): too many images for fixed engine batch size " + + std::to_string(mConfig.batchSize)); + } + formatImage(image, batchIndex, stream); + ++batchIndex; + } + } + + if (batchIndex == 0) + { + return; + } + + if (batchIndex != mConfig.batchSize) + { + throw std::runtime_error("VitRunner::imagePreprocess(): expected " + std::to_string(mConfig.batchSize) + + " images, got " + std::to_string(batchIndex)); + } + + mMultimodalMetrics.recordRun(batchIndex, mConfig.numImageTokens); +} + +void VitRunner::textPreprocess(rt::LLMGenerationRequest const& request, + std::vector>& batchedInputIds, tokenizer::Tokenizer const* tokenizer) +{ + int32_t nextImageTokenId = mConfig.vocabSize; + int64_t numExpandedImageTokens{0}; + + for (size_t i = 0; i < request.requests.size(); ++i) + { + std::vector ids = tokenizer->encode(request.formattedRequests[i].formattedCompleteRequest); + check::check(!ids.empty(), "VitRunner::textPreprocess(): failed to encode text"); + + std::vector newIds; + newIds.reserve(ids.size()); + for (int32_t const id : ids) + { + if (id == mConfig.imageTokenId) + { + for (int32_t k = 0; k < mConfig.seqLenPerImage; ++k) + { + newIds.push_back(nextImageTokenId); + ++nextImageTokenId; + ++numExpandedImageTokens; + } + } + else + { + newIds.push_back(id); + } + } + batchedInputIds.emplace_back(std::move(newIds)); + } + + check::check(numExpandedImageTokens == mConfig.numImageTokens, + "VitRunner::textPreprocess(): expanded image token count must match engine output rows"); +} + +bool VitRunner::preprocess(rt::LLMGenerationRequest const& request, std::vector>& batchedInputIds, + tokenizer::Tokenizer const* tokenizer, rt::OptionalOutputTensor /*mropeCosSinOut*/, cudaStream_t stream, + bool imageOnly) noexcept +{ + try + { + imagePreprocess(request, stream); + // VitRunner uses standard RoPE, so mropeCosSinOut is unused. imageOnly skips text tokenization + // for benchmarking paths that only need the visual engine inputs. + if (!imageOnly) + { + textPreprocess(request, batchedInputIds, tokenizer); + } + } + catch (std::exception const& e) + { + LOG_ERROR("VitRunner::preprocess failed: %s", e.what()); + return false; + } + + return true; +} + +bool VitRunner::infer(cudaStream_t stream) +{ + if (mVitInput.getShape().volume() == 0) + { + return true; + } + + TIME_STAGE(metrics::StageNames::kMULTIMODAL_PROCESSING, stream); + + if (!mVisualContext->setInputShape(binding_names::kVitInput, mVitInput.getShape().getTRTDims())) + { + LOG_ERROR("VitRunner::infer(): failed to set engine input shape"); + return false; + } + + if (!mVisualContext->enqueueV3(stream)) + { + LOG_ERROR("VitRunner::infer(): failed to enqueue engine"); + return false; + } + + return true; +} + +} // namespace rt +} // namespace trt_edgellm diff --git a/cpp/multimodal/vitRunner.h b/cpp/multimodal/vitRunner.h new file mode 100644 index 00000000..0ba08606 --- /dev/null +++ b/cpp/multimodal/vitRunner.h @@ -0,0 +1,85 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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. + */ + +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "multimodal/multimodalRunner.h" +#include "runtime/imageUtils.h" +#include "tokenizer/tokenizer.h" + +#include +#include +#include + +namespace trt_edgellm +{ +namespace rt +{ + +//! Configuration for model-agnostic fixed-shape ViT vision engines. +struct VitConfig +{ + int32_t batchSize{1}; + int32_t numChannels{3}; + int32_t imageHeight{0}; + int32_t imageWidth{0}; + int32_t seqLenPerImage{0}; + int32_t numImageTokens{0}; + int32_t outHiddenSize{0}; + int32_t imageTokenId{0}; + int32_t vocabSize{0}; + std::array imageMean{{0.5F, 0.5F, 0.5F}}; + std::array imageStd{{0.5F, 0.5F, 0.5F}}; +}; + +class VitRunner : public MultimodalRunner +{ +public: + VitRunner(std::string const& engineDir, cudaStream_t stream); + ~VitRunner() noexcept override = default; + + bool validateAndFillConfig(std::string const& configPath) override; + bool allocateBuffer(cudaStream_t stream) override; + + bool preprocess(rt::LLMGenerationRequest const& request, std::vector>& batchedInputIds, + tokenizer::Tokenizer const* tokenizer, rt::OptionalOutputTensor mropeCosSinOut, cudaStream_t stream, + bool imageOnly = false) noexcept override; + + bool infer(cudaStream_t stream) override; + +private: + void imagePreprocess(rt::LLMGenerationRequest const& request, cudaStream_t stream); + void textPreprocess(rt::LLMGenerationRequest const& request, std::vector>& batchedInputIds, + tokenizer::Tokenizer const* tokenizer); + void formatImage(rt::imageUtils::ImageData const& image, int64_t batchIndex, cudaStream_t stream); + + VitConfig mConfig{}; + rt::Tensor mVitInput{}; + rt::Tensor mImageMean{}; + rt::Tensor mImageStd{}; + rt::Tensor mImageDevice{}; + rt::Tensor mNormalizedImageDevice{}; + rt::imageUtils::ImageData mResizedImageHost{}; +}; + +} // namespace rt +} // namespace trt_edgellm diff --git a/cpp/plugins/attentionPlugin/attentionPlugin.cpp b/cpp/plugins/attentionPlugin/attentionPlugin.cpp index 1163c886..c2d45263 100644 --- a/cpp/plugins/attentionPlugin/attentionPlugin.cpp +++ b/cpp/plugins/attentionPlugin/attentionPlugin.cpp @@ -17,7 +17,6 @@ #include "attentionPlugin.h" -#include "common/checkMacros.h" #include "common/cudaUtils.h" #include "common/logger.h" #include "common/tensor.h" @@ -32,11 +31,6 @@ #include "kernels/contextAttentionKernels/cuteDslFMHARunner.h" #endif -// CuTe DSL FFPA kernel (headDim=512 causal attention) -#ifdef CUTE_DSL_FFPA_ENABLED -#include "kernels/contextAttentionKernels/cuteDslFFPARunner.h" -#endif - #include #include #include @@ -57,6 +51,10 @@ namespace constexpr char const* kATTENTION_PLUGIN_VERSION{"1"}; constexpr char const* kATTENTION_PLUGIN_NAME{"AttentionPlugin"}; +// Serialization tag that distinguishes the newer [tag][context_attention_mask_type] trailing payload +// from the legacy single-int32 bidirectional-prefill flag written by older engines. +constexpr int32_t kContextMaskSerializationTag{0x4D41534B}; // 'MASK' + // Select KV cache storage datatype based on FP8 enablement static inline DataType selectKvCacheDataType(bool enableFp8KVCache) { @@ -150,80 +148,6 @@ AttentionExecutionMode deduceModeTreeAttention( return AttentionExecutionMode::kINVALID; } -bool loadFMHAKernels(bool& useCuteDslFMHA, int32_t headSize, int32_t smVersion, nvinfer1::DataType dataType) -{ - bool canImplementFMHA = false; -#ifdef CUTE_DSL_FMHA_ENABLED - if (useCuteDslFMHA) - { - if (CuteDslFMHARunner::canImplement(headSize, smVersion) && CuteDslFMHARunner::loadLLMKernelModule()) - { - canImplementFMHA = true; - LOG_DEBUG("CuTe DSL FMHA kernel loaded for SM%d", smVersion); - } - else - { - LOG_DEBUG("CuTe DSL FMHA not available (headSize=%d, SM%d), falling back to FMHA_v2", headSize, smVersion); - useCuteDslFMHA = false; - } - } - if (!useCuteDslFMHA) -#endif - { - canImplementFMHA = ContextFMHARunner::canImplement( - headSize, smVersion, dataType, AttentionInputLayout::SEPARATE_Q_K_V, ContextAttentionMaskType::CAUSAL); - if (canImplementFMHA) - { - if (!ContextFMHARunner::loadContextFMHAKernels(smVersion, dataType)) - { - LOG_ERROR("Failed to load FMHA_v2 cubins for SM%d", smVersion); - canImplementFMHA = false; - } - } - } - return canImplementFMHA; -} - -// Workspace layout (cumulative, worst-case across all execution paths): -// -// Slot | Shape | Type | Used by -// ------+----------------------------------+-------+------------------------------------------ -// 0 | [B+1] | INT32 | cuQSeqLens (prefill) -// 1 | [B+1] | INT32 | cuKVSeqLens (prefill) -// 2 | [B] | INT32 | kvCacheEndIdxs (prefill) -// 3 | [B+1] | INT32 | paddedCuKVSeqLens (prefill, CuTe DSL) -// 4 | [B, 2, Hkv, Smax, D] | HALF | transposedKV (FMHA_v2 chunked prefill) -// 5* | [B, S, Hq, D] | FP8 | fp8Q (CuTe DSL + FP8 prefill only) -// -// * Slot 5 is conditionally allocated (CuTe DSL + FP8 KV cache only). -// -// Total allocation is the sum of all conditional slots (safe upper bound). -size_t getAttentionWorkspaceSize(int64_t batchSize, int64_t seqLen, int64_t kvCacheCapacity, int32_t numQHeads, - int32_t numKVHeads, int32_t headSize, bool useCuteDslFMHA, bool enableFp8KVCache) -{ - size_t workspaceSize = 0; - - // CuQSeqLens for FMHA. - workspaceSize = accumulateWorkspaceSize(workspaceSize, {batchSize + 1}, DataType::kINT32); - - // Always reserve workspace memory to prepare for chunked prefill decoding. The implementation should be further - // optimized to avoid the workspace size overhead. - workspaceSize = accumulateWorkspaceSize(workspaceSize, rt::Coords{batchSize + 1}, DataType::kINT32); - workspaceSize = accumulateWorkspaceSize(workspaceSize, rt::Coords{batchSize}, DataType::kINT32); - workspaceSize = accumulateWorkspaceSize(workspaceSize, rt::Coords{batchSize + 1}, DataType::kINT32); - workspaceSize = accumulateWorkspaceSize( - workspaceSize, rt::Coords{batchSize, 2, numKVHeads, kvCacheCapacity, headSize}, DataType::kHALF); - - // FP8 Q output: RoPE kernel writes FP8 Q to this workspace buffer (CuTe DSL FMHA path). - if (useCuteDslFMHA && enableFp8KVCache) - { - workspaceSize = accumulateWorkspaceSize( - workspaceSize, rt::Coords{batchSize, seqLen, numQHeads, headSize}, DataType::kFP8); - } - - return workspaceSize; -} - } // namespace // Static class fields initialization @@ -232,58 +156,9 @@ std::vector AttentionPluginCreator::mPluginAttributes; REGISTER_TENSORRT_PLUGIN(AttentionPluginCreator); -// TODO: Extend the attention kernel to consume interleaved KV cache directly and remove this extra -// deinterleaveKVCache call (WAR for current kernel limitation). -std::pair AttentionPlugin::deinterleaveKVCache(rt::Tensor const& kvCacheTensor, - std::byte*& workspacePtr, int32_t batchSize, int32_t numKVHeads, int32_t kvCacheCapacity, int32_t headSize, - int32_t seqLen, cudaStream_t stream) -{ - // seqLen == 0 means copy full capacity; otherwise copy only first seqLen tokens (compact). - int32_t const outSeqDim = (seqLen > 0) ? seqLen : kvCacheCapacity; - size_t const halfSize = static_cast(batchSize) * outSeqDim * numKVHeads * headSize; - rt::Tensor kvWorkspaceTensor - = assignTensorFromWorkspace(workspacePtr, {batchSize, 2, numKVHeads, outSeqDim, headSize}, DataType::kHALF); - half* ptr = kvWorkspaceTensor.dataPointer(); - rt::Tensor kTensor( - ptr, rt::Coords{batchSize, outSeqDim, numKVHeads, headSize}, rt::DeviceType::kGPU, DataType::kHALF); - rt::Tensor vTensor( - ptr + halfSize, rt::Coords{batchSize, outSeqDim, numKVHeads, headSize}, rt::DeviceType::kGPU, DataType::kHALF); - - // seqLen > 0: compact copy of first seqLen tokens; seqLen == 0: full copy (also handles FP8 dequant). - kernel::cvtKVLayoutBHSDToSplitKV(kvCacheTensor, kTensor, vTensor, rt::Tensor{}, seqLen, stream); - return std::make_pair(std::move(kTensor), std::move(vTensor)); -} - -#ifdef CUTE_DSL_FFPA_ENABLED -void AttentionPlugin::dispatchFFPAKernel(half const* q, half const* k, half const* v, half* o, int32_t batchSize, - int32_t seqlenQ, int32_t seqlenK, int32_t numQHeads, int32_t numKVHeads, int32_t headDim, cudaStream_t stream) -{ - CuteDslFFPAParams ffpaParams{}; - ffpaParams.q = q; - ffpaParams.k = k; - ffpaParams.v = v; - ffpaParams.o = o; - ffpaParams.batchSize = batchSize; - ffpaParams.seqlenQ = seqlenQ; - ffpaParams.seqlenK = seqlenK; - ffpaParams.numQHeads = numQHeads; - ffpaParams.numKVHeads = numKVHeads; - ffpaParams.headDim = headDim; - ffpaParams.softmaxScale = 1.0F / std::sqrt(static_cast(headDim)); - CuteDslFFPARunner::run(ffpaParams, stream); -} -#endif - -void AttentionPlugin::zeroPrefillOutputForPaddingForFFPA(rt::Tensor& attentionOutput, int32_t batchSize, int32_t seqLen, - int32_t numQHeads, int32_t headSize, cudaStream_t stream) -{ - size_t const outputBytes = static_cast(batchSize) * seqLen * numQHeads * headSize * sizeof(half); - CUDA_CHECK(cudaMemsetAsync(attentionOutput.rawPointer(), 0, outputBytes, stream)); -} - AttentionPlugin::AttentionPlugin(std::string const& name, int32_t numQHeads, int32_t numKVHeads, int32_t headSize, int32_t enableTreeAttention, int32_t enableFp8KVCache, int32_t slidingWindowSize, - std::vector const& qkvScales) + std::vector const& qkvScales, int32_t contextAttentionMaskType) : mLayerName(name) , mNumQHeads(numQHeads) , mNumKVHeads(numKVHeads) @@ -292,11 +167,14 @@ AttentionPlugin::AttentionPlugin(std::string const& name, int32_t numQHeads, int , mEnableFp8KVCache(enableFp8KVCache) , mQkvScales(enableFp8KVCache ? qkvScales : std::vector{1.f, 1.f, 1.f}) , mSlidingWindowSize(slidingWindowSize) + , mContextAttentionMaskType(contextAttentionMaskType) { - ELLM_CHECK(!mEnableFp8KVCache || mQkvScales.size() == 3, - "FP8 KV cache enabled but qkv_scales has " + if (mEnableFp8KVCache && mQkvScales.size() != 3) + { + throw std::runtime_error("FP8 KV cache enabled but qkv_scales has " + std::to_string(mQkvScales.size()) + " elements (expected 3). " "Re-export the model to include QKV scales [q, k, v]."); + } mSMVersion = getSMVersion(); applyThorSMRenumberWAR(mSMVersion); @@ -304,182 +182,166 @@ AttentionPlugin::AttentionPlugin(std::string const& name, int32_t numQHeads, int LOG_DEBUG("AttentionPlugin FMHA path: %s, sliding_window: %s", mUseCuteDslFMHA ? "CuTe DSL FMHA" : "FMHA_v2", mSlidingWindowSize > 0 ? std::to_string(mSlidingWindowSize).c_str() : "disabled"); - mCanImplementFMHA = loadFMHAKernels(mUseCuteDslFMHA, mHeadSize, mSMVersion, mDataType); + auto const prefillMaskType = static_cast(mContextAttentionMaskType); + + // Check FMHA implementation support and load the corresponding kernel module. + bool canImplementFMHA = false; +#ifdef CUTE_DSL_FMHA_ENABLED + // CuTe DSL FMHA currently supports the causal prefill path only. + bool const useCausalFMHA = prefillMaskType == ContextAttentionMaskType::CAUSAL; + if (useCausalFMHA && mUseCuteDslFMHA && CuteDslFMHARunner::canImplement(mHeadSize, mSMVersion)) + { + if (CuteDslFMHARunner::loadLLMKernelModule()) + { + canImplementFMHA = true; + LOG_DEBUG("CuTe DSL FMHA kernel loaded for SM%d", mSMVersion); + } + else + { + LOG_WARNING("CuTe DSL FMHA kernel failed to load, falling back to FMHA_v2"); + mUseCuteDslFMHA = false; + } + } + if (!canImplementFMHA) +#endif + { + // Fallback to FMHA_v2 cubins. + canImplementFMHA = ContextFMHARunner::canImplement( + mHeadSize, mSMVersion, mDataType, AttentionInputLayout::SEPARATE_Q_K_V, prefillMaskType); + if (canImplementFMHA) + { + if (!ContextFMHARunner::loadContextFMHAKernels(mSMVersion, mDataType)) + { + LOG_ERROR("Failed to load FMHA_v2 cubins for SM%d", mSMVersion); + canImplementFMHA = false; + } + } + } - // XQA decode kernels are needed for decode path when available. - bool const useSpecDecode = true; - mCanImplementXQA = DecoderXQARunner::canImplement( + // XQA decode kernels are always needed regardless of FMHA path. + bool const useSpecDecode = static_cast(mEnableTreeAttention); + bool canImplementXQA = DecoderXQARunner::canImplement( mNumQHeads, mNumKVHeads, mHeadSize, mSMVersion, mDataType, selectKvCacheDataType(mEnableFp8KVCache)); - if (mCanImplementXQA) + if (canImplementXQA) { DecoderXQARunner::loadDecodeXQAKernels( mSMVersion, mDataType, selectKvCacheDataType(mEnableFp8KVCache), useSpecDecode); } - // Kernel selection priority for prefill and decode: - // 1. FMHA (prefill) + XQA (decode) — standard path for most head sizes. - // 2. FFPA (prefill) + XQA (decode) — fallback for headSize=512 where FMHA has no cubins. - // 3. FFPA (prefill) only — headSize=512 without XQA decode support. - // 4. XQA (decode) only — naive attention for prefill (degraded). - // 5. None — fatal, cannot serve this configuration. - if (mCanImplementFMHA) + if (!canImplementFMHA || !canImplementXQA) { - LOG_INFO("AttentionPlugin: FMHA supported for headSize=%d, using FMHA for prefill%s.", mHeadSize, - mCanImplementXQA ? " + XQA for decode" : ""); + LOG_ERROR( + "Cannot implement AttentionPlugin configuration. FMHA: %s, XQA: %s, SM: %d, HeadSize: %d, NumQHeads: %d, " + "NumKVHeads: %d", + canImplementFMHA ? "supported" : "NOT supported", canImplementXQA ? "supported" : "NOT supported", + mSMVersion, mHeadSize, mNumQHeads, mNumKVHeads); + throw std::runtime_error("Cannot implement the AttentionPlugin configuration."); } - else +} + +AttentionPlugin::AttentionPlugin(std::string const& name, std::byte const* data, size_t length) + : mLayerName(name) +{ + deserializeValue(&data, &length, &mNumQHeads); + deserializeValue(&data, &length, &mNumKVHeads); + deserializeValue(&data, &length, &mHeadSize); + deserializeValue(&data, &length, &mEnableTreeAttention); + deserializeValue(&data, &length, &mEnableFp8KVCache); + deserializeValue(&data, &length, &mSlidingWindowSize); + + // Optional trailing fields: [int32 qkv_scale_count][float * count] followed by either the newer + // [int32 tag][int32 context_attention_mask_type] payload or the legacy [int32 bidirectional_prefill] + // flag. Engines serialized before these fields default to unit QKV scales and causal prefill. + if (length >= sizeof(int32_t)) { - // FMHA unavailable — try FFPA d512 kernel as prefill fallback for headSize=512. -#ifdef CUTE_DSL_FFPA_ENABLED - if (mHeadSize == 512 && CuteDslFFPARunner::canImplement(mHeadSize, mSMVersion)) + int32_t qkvScaleCount = 0; + deserializeValue(&data, &length, &qkvScaleCount); + if (qkvScaleCount < 0 || length < static_cast(qkvScaleCount) * sizeof(float)) { - if (CuteDslFFPARunner::loadKernelModule()) - { - mCanImplementFFPA = true; - } - else - { - LOG_WARNING("AttentionPlugin: Failed to load FFPA d512 kernel."); - } + throw std::runtime_error("Invalid qkv_scales payload in serialized AttentionPlugin data."); } -#endif - if (mCanImplementFFPA && mCanImplementXQA) + std::vector qkvScales(qkvScaleCount); + for (int32_t i = 0; i < qkvScaleCount; ++i) { - LOG_INFO("AttentionPlugin: FMHA unsupported for headSize=%d, using FFPA for prefill + XQA for decode.", - mHeadSize); + deserializeValue(&data, &length, &qkvScales[i]); } - else if (mCanImplementFFPA) + if (qkvScaleCount == 3) { - LOG_INFO("AttentionPlugin: FMHA/XQA unsupported for headSize=%d numKVHeads=%d, using FFPA for prefill.", - mHeadSize, mNumKVHeads); - } - else if (mCanImplementXQA) - { - LOG_WARNING( - "AttentionPlugin: FMHA/FFPA unsupported for headSize=%d, using naive attention for prefill + XQA for " - "decode.", - mHeadSize); - } - else - { - LOG_ERROR( - "Cannot implement AttentionPlugin configuration. FMHA: %s, XQA: %s, FFPA: %s, SM: %d, HeadSize: %d, " - "NumQHeads: %d, NumKVHeads: %d", - "NOT supported", "NOT supported", "NOT supported", mSMVersion, mHeadSize, mNumQHeads, mNumKVHeads); - throw std::runtime_error("Cannot implement the AttentionPlugin configuration."); + mQkvScales = qkvScales; } } -} -AttentionPlugin::AttentionPlugin(std::string const& name, PluginFieldCollection const* fc) - : mLayerName(name) -{ - mNumQHeads = parsePluginScalarField("num_q_heads", fc).value_or(0); - mNumKVHeads = parsePluginScalarField("num_kv_heads", fc).value_or(0); - mHeadSize = parsePluginScalarField("head_size", fc).value_or(0); - mEnableTreeAttention = parsePluginScalarField("enable_tree_attention", fc).value_or(0); - mEnableFp8KVCache = parsePluginScalarField("enable_fp8_kv_cache", fc).value_or(0); - mSlidingWindowSize = parsePluginScalarField("sliding_window_size", fc).value_or(-1); - - // Parse qkv_scales float array - for (int32_t i = 0; i < fc->nbFields; ++i) + if (mEnableFp8KVCache && mQkvScales.size() != 3) { - if (std::string("qkv_scales") == fc->fields[i].name) - { - auto const* data = static_cast(fc->fields[i].data); - mQkvScales.assign(data, data + fc->fields[i].length); - break; - } + throw std::runtime_error( + "FP8 KV cache enabled but qkv_scales missing or incomplete in serialized AttentionPlugin data."); } - if (!mEnableFp8KVCache) + // New engines write [int32 tag][int32 maskType]; legacy engines wrote a single [int32 + // bidirectionalPrefill] flag (0 = causal, 1 = bidirectional/padding). Migrate the legacy flag. + if (length >= 2 * sizeof(int32_t)) { - mQkvScales = {1.f, 1.f, 1.f}; + int32_t maskTag = 0; + deserializeValue(&data, &length, &maskTag); + check::check(maskTag == kContextMaskSerializationTag, + "Unexpected serialization tag for AttentionPlugin context attention mask type"); + deserializeValue(&data, &length, &mContextAttentionMaskType); } - else + else if (length >= sizeof(int32_t)) { - ELLM_CHECK(mQkvScales.size() == 3, - "FP8 KV cache enabled but qkv_scales missing or incomplete " - "in plugin fields (expected 3). Re-export the model with QKV scales [q, k, v]."); + int32_t legacyBidirectionalPrefill = 0; + deserializeValue(&data, &length, &legacyBidirectionalPrefill); + mContextAttentionMaskType = (legacyBidirectionalPrefill != 0) + ? static_cast(ContextAttentionMaskType::PADDING) + : static_cast(ContextAttentionMaskType::CAUSAL); } + check::check(length == 0, "Unexpected trailing bytes in serialized AttentionPlugin data"); mSMVersion = getSMVersion(); applyThorSMRenumberWAR(mSMVersion); LOG_DEBUG("AttentionPlugin FMHA path: %s", mUseCuteDslFMHA ? "CuTe DSL FMHA" : "FMHA_v2"); - mCanImplementFMHA = loadFMHAKernels(mUseCuteDslFMHA, mHeadSize, mSMVersion, mDataType); - - // XQA decode kernels. - mCanImplementXQA = DecoderXQARunner::canImplement( - mNumQHeads, mNumKVHeads, mHeadSize, mSMVersion, mDataType, selectKvCacheDataType(mEnableFp8KVCache)); - if (mCanImplementXQA) - { - DecoderXQARunner::loadDecodeXQAKernels( - mSMVersion, mDataType, selectKvCacheDataType(mEnableFp8KVCache), /*useSpecDecodeKernels=*/true); - } - - if (!mCanImplementFMHA) + // Load FMHA kernel module based on implementation support. +#ifdef CUTE_DSL_FMHA_ENABLED + // CuTe DSL FMHA currently supports the causal prefill path only. + bool const useCausalFMHA + = static_cast(mContextAttentionMaskType) == ContextAttentionMaskType::CAUSAL; + if (useCausalFMHA && mUseCuteDslFMHA && CuteDslFMHARunner::canImplement(mHeadSize, mSMVersion)) { -#ifdef CUTE_DSL_FFPA_ENABLED - if (mHeadSize == 512 && CuteDslFFPARunner::canImplement(mHeadSize, mSMVersion)) + if (!CuteDslFMHARunner::loadLLMKernelModule()) { - if (CuteDslFFPARunner::loadKernelModule()) - { - mCanImplementFFPA = true; - } + LOG_WARNING("CuTe DSL FMHA kernel failed to load, falling back to FMHA_v2"); + mUseCuteDslFMHA = false; } -#endif } -} - -AttentionPlugin::~AttentionPlugin() = default; - -// --------------------------------------------------------------------------- -// IPluginV3 -// --------------------------------------------------------------------------- - -IPluginCapability* AttentionPlugin::getCapabilityInterface(PluginCapabilityType type) noexcept -{ - try + if (!useCausalFMHA || !mUseCuteDslFMHA) +#endif { - if (type == PluginCapabilityType::kBUILD) - { - return static_cast(this); - } - if (type == PluginCapabilityType::kRUNTIME) + if (!ContextFMHARunner::loadContextFMHAKernels(mSMVersion, mDataType)) { - return static_cast(this); + LOG_ERROR("Failed to load FMHA_v2 cubins for SM%d", mSMVersion); } - return static_cast(this); - } - catch (std::exception const& e) - { - return nullptr; } + + // XQA decode kernels are always needed regardless of FMHA path. + bool const useSpecDecode = static_cast(mEnableTreeAttention); + DecoderXQARunner::loadDecodeXQAKernels( + mSMVersion, mDataType, selectKvCacheDataType(mEnableFp8KVCache), useSpecDecode); } -IPluginV3* AttentionPlugin::clone() noexcept +AttentionPlugin::~AttentionPlugin() {} + +IPluginV2DynamicExt* AttentionPlugin::clone() const noexcept { - try - { - auto* p = new AttentionPlugin(mLayerName, mNumQHeads, mNumKVHeads, mHeadSize, mEnableTreeAttention, - mEnableFp8KVCache, mSlidingWindowSize, mQkvScales); - p->setPluginNamespace(mNamespace.c_str()); - return p; - } - catch (...) - { - return nullptr; - } + AttentionPlugin* plugin = new AttentionPlugin(mLayerName, mNumQHeads, mNumKVHeads, mHeadSize, mEnableTreeAttention, + mEnableFp8KVCache, mSlidingWindowSize, mQkvScales, mContextAttentionMaskType); + plugin->setPluginNamespace(mNamespace.c_str()); + return plugin; } -// --------------------------------------------------------------------------- -// IPluginV3OneCore — metadata -// --------------------------------------------------------------------------- - -char const* AttentionPlugin::getPluginName() const noexcept +char const* AttentionPlugin::getPluginType() const noexcept { return kATTENTION_PLUGIN_NAME; } @@ -491,7 +353,7 @@ char const* AttentionPlugin::getPluginNamespace() const noexcept void AttentionPlugin::setPluginNamespace(char const* pluginNamespace) noexcept { - mNamespace = pluginNamespace ? pluginNamespace : ""; + mNamespace = std::string(pluginNamespace); } char const* AttentionPlugin::getPluginVersion() const noexcept @@ -499,66 +361,14 @@ char const* AttentionPlugin::getPluginVersion() const noexcept return kATTENTION_PLUGIN_VERSION; } -// --------------------------------------------------------------------------- -// IPluginV3OneBuild — shape / format -// --------------------------------------------------------------------------- - int32_t AttentionPlugin::getNbOutputs() const noexcept { // At both context and generation phase, output attention result and kv-cache. return 2; } -int32_t AttentionPlugin::getOutputDataTypes(DataType* outputTypes, [[maybe_unused]] int32_t nbOutputs, - DataType const* inputTypes, [[maybe_unused]] int32_t nbInputs) const noexcept -{ - try - { - assert(nbOutputs == kNUM_REQUIRED_OUTPUTS); - // Output[0] (attention): always FP16 (follows Q input dtype). - // Output[1] (KV cache) follows KV input dtype (HALF or FP8). - outputTypes[kOUT_ATTENTION_IDX] = inputTypes[kIN_Q_IDX]; - outputTypes[kOUT_KV_CACHE_IDX] = inputTypes[kIN_KV_CACHE_IDX]; - return 0; - } - catch (std::exception const& e) - { - return -1; - } -} - -int32_t AttentionPlugin::getOutputShapes(DimsExprs const* inputs, [[maybe_unused]] int32_t nbInputs, - DimsExprs const* /* shapeInputs */, int32_t /* nbShapeInputs */, DimsExprs* outputs, - [[maybe_unused]] int32_t nbOutputs, IExprBuilder& exprBuilder) noexcept -{ - try - { - assert(nbOutputs == kNUM_REQUIRED_OUTPUTS); - // Output[0] is attention result, has shape [B, S, Hq, D]. Refers to Q shape [B, S, Hq*D] - outputs[kOUT_ATTENTION_IDX].nbDims = 4; - outputs[kOUT_ATTENTION_IDX].d[0] = inputs[kIN_Q_IDX].d[0]; - outputs[kOUT_ATTENTION_IDX].d[1] = inputs[kIN_Q_IDX].d[1]; - outputs[kOUT_ATTENTION_IDX].d[2] = exprBuilder.constant(mNumQHeads); - outputs[kOUT_ATTENTION_IDX].d[3] = exprBuilder.constant(mHeadSize); - - // Output[1] is KVCache, same shape as input KV cache [B, 2, Hkv, Smax, D] - outputs[kOUT_KV_CACHE_IDX].nbDims = 5; - outputs[kOUT_KV_CACHE_IDX].d[0] = inputs[kIN_KV_CACHE_IDX].d[0]; - outputs[kOUT_KV_CACHE_IDX].d[1] = inputs[kIN_KV_CACHE_IDX].d[1]; - outputs[kOUT_KV_CACHE_IDX].d[2] = inputs[kIN_KV_CACHE_IDX].d[2]; - outputs[kOUT_KV_CACHE_IDX].d[3] = inputs[kIN_KV_CACHE_IDX].d[3]; - outputs[kOUT_KV_CACHE_IDX].d[4] = inputs[kIN_KV_CACHE_IDX].d[4]; - - return 0; - } - catch (std::exception const& e) - { - return -1; - } -} - bool AttentionPlugin::supportsFormatCombination( - int32_t pos, DynamicPluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept + int32_t pos, nvinfer1::PluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept { // Support context/generation phase inputs: // Q tensor (linear FP16) with shape [B, S, Hq, D] @@ -577,28 +387,33 @@ bool AttentionPlugin::supportsFormatCombination( // Support context/generation phase outputs: // attention result (linear FP16) with shape [B, S, Hq, D] // KV-cache tensor, same as the above. - // NOTE: Q/K/V/KVCache dimension-value assertions (e.g. d[2] == mNumQHeads * mHeadSize) - // are intentionally omitted here to support Gemma4's heterogeneous per-layer head - // configurations (shared-KV layers have d[2]=0 for K/V). Full shape validation is - // performed at runtime in enqueue() where actual tensor dimensions are checked against - // mNumQHeads, mNumKVHeads, and mHeadSize. - auto checkQ = [](PluginTensorDesc const& tensorDesc) { + auto checkQ = [this](nvinfer1::PluginTensorDesc const& tensorDesc) { bool status{true}; status &= tensorDesc.type == DataType::kHALF; status &= tensorDesc.format == TensorFormat::kLINEAR; status &= tensorDesc.dims.nbDims == 3; + auto const tensorDim = tensorDesc.dims; + if (status) + { + status &= tensorDim.d[2] == mNumQHeads * mHeadSize; + } return status; }; - auto checkKV = [](PluginTensorDesc const& tensorDesc) { + auto checkKV = [this](nvinfer1::PluginTensorDesc const& tensorDesc) { bool status{true}; status &= tensorDesc.type == DataType::kHALF; status &= tensorDesc.format == TensorFormat::kLINEAR; status &= tensorDesc.dims.nbDims == 3; + auto const tensorDim = tensorDesc.dims; + if (status) + { + status &= tensorDim.d[2] == mNumKVHeads * mHeadSize; + } return status; }; - auto checkKVCache = [this](PluginTensorDesc const& tensorDesc) { + auto checkKVCache = [this](nvinfer1::PluginTensorDesc const& tensorDesc) { bool status{true}; // Support FP16 or FP8 storage; if (mEnableFp8KVCache) @@ -611,10 +426,17 @@ bool AttentionPlugin::supportsFormatCombination( } status &= tensorDesc.format == TensorFormat::kLINEAR; status &= tensorDesc.dims.nbDims == 5; + if (status) + { + auto const tensorDim = tensorDesc.dims; + status &= tensorDim.d[1] == 2; // Specify K and V + status &= tensorDim.d[2] == mNumKVHeads; + status &= tensorDim.d[4] == mHeadSize; + } return status; }; - auto checkSequenceLen = [](PluginTensorDesc const& tensorDesc) { + auto checkSequenceLen = [this](nvinfer1::PluginTensorDesc const& tensorDesc) { bool status{true}; status &= tensorDesc.type == DataType::kINT32; status &= tensorDesc.format == TensorFormat::kLINEAR; @@ -622,7 +444,7 @@ bool AttentionPlugin::supportsFormatCombination( return status; }; - auto checkPosEncodingCosSin = [this](PluginTensorDesc const& tensorDesc) { + auto checkPosEncodingCosSin = [this](nvinfer1::PluginTensorDesc const& tensorDesc) { bool status{true}; status &= tensorDesc.type == DataType::kFLOAT; status &= tensorDesc.format == TensorFormat::kLINEAR; @@ -631,7 +453,7 @@ bool AttentionPlugin::supportsFormatCombination( return status; }; - auto checkAttentionMask = [](PluginTensorDesc const& tensorDesc) { + auto checkAttentionMask = [this](nvinfer1::PluginTensorDesc const& tensorDesc) { bool status{true}; status &= tensorDesc.type == DataType::kINT32; status &= tensorDesc.format == TensorFormat::kLINEAR; @@ -639,7 +461,7 @@ bool AttentionPlugin::supportsFormatCombination( return status; }; - auto checkAttentionPosId = [](PluginTensorDesc const& tensorDesc) { + auto checkAttentionPosId = [this](nvinfer1::PluginTensorDesc const& tensorDesc) { bool status{true}; status &= tensorDesc.type == DataType::kINT32; status &= tensorDesc.format == TensorFormat::kLINEAR; @@ -647,7 +469,7 @@ bool AttentionPlugin::supportsFormatCombination( return status; }; - auto checkKVCacheStartIdx = [](PluginTensorDesc const& tensorDesc) { + auto checkKVCacheStartIdx = [this](nvinfer1::PluginTensorDesc const& tensorDesc) { bool status{true}; status &= tensorDesc.type == DataType::kINT32; status &= tensorDesc.format == TensorFormat::kLINEAR; @@ -655,11 +477,17 @@ bool AttentionPlugin::supportsFormatCombination( return status; }; - auto checkAttentionOutput = [](PluginTensorDesc const& tensorDesc) { + auto checkAttentionOutput = [this](nvinfer1::PluginTensorDesc const& tensorDesc) { bool status{true}; status &= tensorDesc.type == DataType::kHALF; status &= tensorDesc.format == TensorFormat::kLINEAR; status &= tensorDesc.dims.nbDims == 4; + if (status) + { + auto const tensorDim = tensorDesc.dims; + status &= tensorDim.d[2] == mNumQHeads; + status &= tensorDim.d[3] == mHeadSize; + } return status; }; @@ -680,13 +508,13 @@ bool AttentionPlugin::supportsFormatCombination( { switch (pos) { - case kIN_Q_IDX: result = checkQ(inOut[pos].desc); break; - case kIN_K_IDX: result = checkKV(inOut[pos].desc); break; - case kIN_V_IDX: result = checkKV(inOut[pos].desc); break; - case kIN_KV_CACHE_IDX: result = checkKVCache(inOut[pos].desc); break; - case kIN_CONTEXT_LENGTH_IDX: result = checkSequenceLen(inOut[pos].desc); break; - case kIN_ROPE_COS_SIN_IDX: result = checkPosEncodingCosSin(inOut[pos].desc); break; - case kIN_KV_CACHE_START_IDX: result = checkKVCacheStartIdx(inOut[pos].desc); break; + case kIN_Q_IDX: result = checkQ(inOut[pos]); break; + case kIN_K_IDX: result = checkKV(inOut[pos]); break; + case kIN_V_IDX: result = checkKV(inOut[pos]); break; + case kIN_KV_CACHE_IDX: result = checkKVCache(inOut[pos]); break; + case kIN_CONTEXT_LENGTH_IDX: result = checkSequenceLen(inOut[pos]); break; + case kIN_ROPE_COS_SIN_IDX: result = checkPosEncodingCosSin(inOut[pos]); break; + case kIN_KV_CACHE_START_IDX: result = checkKVCacheStartIdx(inOut[pos]); break; default: break; } @@ -698,12 +526,12 @@ bool AttentionPlugin::supportsFormatCombination( { if (pos == currentOptionalInputIdx) { - result = checkAttentionMask(inOut[pos].desc); + result = checkAttentionMask(inOut[pos]); } currentOptionalInputIdx++; if (pos == currentOptionalInputIdx) { - result = checkAttentionPosId(inOut[pos].desc); + result = checkAttentionPosId(inOut[pos]); } currentOptionalInputIdx++; } @@ -714,8 +542,8 @@ bool AttentionPlugin::supportsFormatCombination( int32_t outPos = pos - nbInputs; switch (outPos) { - case kOUT_ATTENTION_IDX: result = checkAttentionOutput(inOut[pos].desc); break; - case kOUT_KV_CACHE_IDX: result = checkKVCache(inOut[pos].desc); break; + case kOUT_ATTENTION_IDX: result = checkAttentionOutput(inOut[pos]); break; + case kOUT_KV_CACHE_IDX: result = checkKVCache(inOut[pos]); break; default: break; } } @@ -723,46 +551,122 @@ bool AttentionPlugin::supportsFormatCombination( return result; } -int32_t AttentionPlugin::configurePlugin([[maybe_unused]] DynamicPluginTensorDesc const* in, - [[maybe_unused]] int32_t nbInputs, [[maybe_unused]] DynamicPluginTensorDesc const* out, +// IPluginV2Ext Methods +DataType AttentionPlugin::getOutputDataType([[maybe_unused]] int32_t index, + [[maybe_unused]] nvinfer1::DataType const* inputTypes, [[maybe_unused]] int32_t nbInputs) const noexcept +{ + // Output[0] (attention): always FP16 (follows Q input dtype). + // Output[1] (KV cache) follows KV input dtype (HALF or FP8). + if (index == kOUT_ATTENTION_IDX) + { + return inputTypes[kIN_Q_IDX]; + } + return inputTypes[kIN_KV_CACHE_IDX]; +} + +DimsExprs AttentionPlugin::getOutputDimensions(int32_t outputIndex, nvinfer1::DimsExprs const* inputs, + [[maybe_unused]] int32_t nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept +{ + // Output[0] is attention result, has shape [B, S, Hq, D]. Refers to Q shape [B, S, Hq, D] + DimsExprs output; + if (outputIndex == kOUT_ATTENTION_IDX) + { + output.nbDims = 4; + output.d[0] = inputs[kIN_Q_IDX].d[0]; + output.d[1] = inputs[kIN_Q_IDX].d[1]; + output.d[2] = exprBuilder.constant(mNumQHeads); + output.d[3] = exprBuilder.constant(mHeadSize); + } + else if (outputIndex == kOUT_KV_CACHE_IDX) + { + // Output[1] is KVCache, same shape as input KV cache + output.nbDims = 5; + output.d[0] = inputs[kIN_KV_CACHE_IDX].d[0]; + output.d[1] = inputs[kIN_KV_CACHE_IDX].d[1]; + output.d[2] = inputs[kIN_KV_CACHE_IDX].d[2]; + output.d[3] = inputs[kIN_KV_CACHE_IDX].d[3]; + output.d[4] = inputs[kIN_KV_CACHE_IDX].d[4]; + } + return output; +} + +void AttentionPlugin::configurePlugin([[maybe_unused]] nvinfer1::DynamicPluginTensorDesc const* in, + [[maybe_unused]] int32_t nbInputs, [[maybe_unused]] nvinfer1::DynamicPluginTensorDesc const* out, [[maybe_unused]] int32_t nbOutputs) noexcept { - return 0; // No need to configure anything since we will only use the runtime tensor shapes. + return; // No need to configure anything since we will only use the runtime tensor shapes. } -size_t AttentionPlugin::getWorkspaceSize(DynamicPluginTensorDesc const* inputs, [[maybe_unused]] int32_t nbInputs, - [[maybe_unused]] DynamicPluginTensorDesc const* outputs, [[maybe_unused]] int32_t nbOutputs) const noexcept +// Workspace layout (cumulative, worst-case across all execution paths): +// +// Slot | Shape | Type | Used by +// ------+----------------------------------+-------+------------------------------------------ +// 0 | [B+1] | INT32 | cuQSeqLens (prefill) +// 1 | [B+1] | INT32 | cuKVSeqLens (prefill) +// 2 | [B] | INT32 | kvCacheEndIdxs (prefill) +// 3 | [B+1] | INT32 | paddedCuKVSeqLens (prefill, CuTe DSL) +// 4 | [B, 2, Hkv, Smax, D] | HALF | transposedKV (FMHA_v2 chunked prefill) +// 5* | [B, S, Hq, D] | FP8 | fp8Q (CuTe DSL + FP8 prefill only) +// +// * Slot 5 is conditionally allocated (CuTe DSL + FP8 KV cache only). +// +// Total allocation is the sum of all conditional slots (safe upper bound). +size_t AttentionPlugin::getWorkspaceSize([[maybe_unused]] nvinfer1::PluginTensorDesc const* inputs, + [[maybe_unused]] int32_t nbInputs, [[maybe_unused]] nvinfer1::PluginTensorDesc const* outputs, + [[maybe_unused]] int32_t nbOutputs) const noexcept { - int64_t const maxBatchSize = inputs[kIN_Q_IDX].max.d[0]; - int64_t const maxSeqLen = inputs[kIN_Q_IDX].max.d[1]; - // KV cache tensor shape: [B, 2, num_kv_heads, capacity, head_dim] - int64_t const maxKVCacheCapacity = inputs[kIN_KV_CACHE_IDX].max.d[3]; - size_t const workspaceSize = getAttentionWorkspaceSize(maxBatchSize, maxSeqLen, maxKVCacheCapacity, mNumQHeads, - mNumKVHeads, mHeadSize, mUseCuteDslFMHA, mEnableFp8KVCache); + // TensorRT will supply max profile shape for each input/output tensor across all optimization profiles. + // We will request workspace to keep intermediate tensors under prefill/decode phase executions. + // Obtain max supported batch size from the input tensor shapes. The Q input tensor will be in shape + // [B, S, Hq, D] where S is padded the max length of the input sequence within this batch. + PluginTensorDesc const& qInputDesc = inputs[kIN_Q_IDX]; + int64_t const maxBatchSize = qInputDesc.dims.d[0]; + + // Obtain max KV cache capacity from the KV cache tensor shape. + // The KV cache tensor has shape [B, 2, num_kv_heads, capacity, head_dim] + PluginTensorDesc const& kvCacheDesc = inputs[kIN_KV_CACHE_IDX]; + int64_t const maxKVCacheCapacity = kvCacheDesc.dims.d[3]; + + size_t workspaceSize = 0; + + // CuQSeqLens for FMHA. + workspaceSize = accumulateWorkspaceSize(workspaceSize, {maxBatchSize + 1}, DataType::kINT32); + + // Always reserve workspace memory to prepare for chunked prefill decoding. The implementation should be further + // optimized to avoid the workspace size overhead. + + // CuTotalKvCacheLens to describe the cumulative length of KV tensors. + workspaceSize = accumulateWorkspaceSize(workspaceSize, rt::Coords{maxBatchSize + 1}, DataType::kINT32); + // KVCache ends that denote the end index of each KVCache lane after adding current contents. + workspaceSize = accumulateWorkspaceSize(workspaceSize, rt::Coords{maxBatchSize}, DataType::kINT32); + // Padded cumulative KV sequence lengths for CuTe DSL FMHA. + workspaceSize = accumulateWorkspaceSize(workspaceSize, rt::Coords{maxBatchSize + 1}, DataType::kINT32); + // KV workspace for split K and V tensors (split into K and V halves by pointer arithmetic in enqueue). + workspaceSize = accumulateWorkspaceSize( + workspaceSize, rt::Coords{maxBatchSize, 2, mNumKVHeads, maxKVCacheCapacity, mHeadSize}, DataType::kHALF); + + // FP8 Q output: RoPE kernel writes FP8 Q to this workspace buffer (CuTe DSL FMHA path). +#ifdef CUTE_DSL_FMHA_ENABLED + if (mUseCuteDslFMHA && mEnableFp8KVCache) + { + int64_t const maxSeqLen = qInputDesc.dims.d[1]; + workspaceSize = accumulateWorkspaceSize( + workspaceSize, rt::Coords{maxBatchSize, maxSeqLen, mNumQHeads, mHeadSize}, DataType::kFP8); + } +#endif + + // Request another alignment size to align the workspace pointer. + workspaceSize += kDEVICE_ALIGNMENT; LOG_DEBUG("AttentionPlugin workspace size: %zu bytes", workspaceSize); return workspaceSize; } -int32_t AttentionPlugin::getAliasedInput(int32_t outputIndex) noexcept +int32_t AttentionPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, + [[maybe_unused]] nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, + void* workspace, cudaStream_t stream) noexcept { - // WAR:this is not the correct plugin API usage. The - // plugin updates the KV cache in place, so the correct return is - // kIN_KV_CACHE_IDX (output kOUT_KV_CACHE_IDX aliases that input). We return -1 - // to drop the alias because declaring it makes Myelin keep a redundant - // per-layer KV copy (the perf regression). In-place read-write still works - // because the runtime binds past and present KV to the same address. TODO: - // restore the alias declaration once the Myelin issue is fixed. - return -1; -} - -// --------------------------------------------------------------------------- -// IPluginV3OneRuntime — execution -// --------------------------------------------------------------------------- -int32_t AttentionPlugin::enqueue(PluginTensorDesc const* inputDesc, [[maybe_unused]] PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept -{ // Construct non-owned tensor objects from I/O data pointers and shapes. // Q input in the graph will be in shape [B, S, Hq x D], for convenience, // we will use shape of [B, S, Hq, D] to represent the tensor. @@ -772,28 +676,20 @@ int32_t AttentionPlugin::enqueue(PluginTensorDesc const* inputDesc, [[maybe_unus PluginTensorDesc const& vInputDesc = inputDesc[kIN_V_IDX]; int32_t const runtimeBatchSize = static_cast(qInputDesc.dims.d[0]); int32_t const runtimeSeqLen = static_cast(qInputDesc.dims.d[1]); - int32_t const kvSeqLen = static_cast(kInputDesc.dims.d[1]); - // Shared-KV layers have no K/V projection: detect by either seq_len=0 or hidden_dim=0. - bool const sharedKV = (kvSeqLen == 0) || (kInputDesc.dims.d[2] == 0); - + check::check(kInputDesc.dims.d[0] == runtimeBatchSize && vInputDesc.dims.d[0] == runtimeBatchSize, + "Batch size must be consistent across Q/K/V inputs."); + check::check(kInputDesc.dims.d[1] == runtimeSeqLen && vInputDesc.dims.d[1] == runtimeSeqLen, + "Sequence length must be consistent across Q/K/V inputs."); check::check(qInputDesc.dims.d[2] == mNumQHeads * mHeadSize, "Q input shape shall be consistent."); - if (!sharedKV) - { - check::check(kInputDesc.dims.d[0] == runtimeBatchSize && vInputDesc.dims.d[0] == runtimeBatchSize, - "Batch size must be consistent across Q/K/V inputs."); - check::check(kInputDesc.dims.d[1] == vInputDesc.dims.d[1], "K and V sequence lengths must be consistent."); - check::check( - kvSeqLen == runtimeSeqLen, "K/V sequence length must equal Q sequence length when not in shared-KV mode."); - check::check(kInputDesc.dims.d[2] == mNumKVHeads * mHeadSize, "K input shape shall be consistent."); - check::check(vInputDesc.dims.d[2] == mNumKVHeads * mHeadSize, "V input shape shall be consistent."); - } + check::check(kInputDesc.dims.d[2] == mNumKVHeads * mHeadSize, "K input shape shall be consistent."); + check::check(vInputDesc.dims.d[2] == mNumKVHeads * mHeadSize, "V input shape shall be consistent."); rt::Tensor qInputTensor(const_cast(inputs[kIN_Q_IDX]), rt::Coords{runtimeBatchSize, runtimeSeqLen, mNumQHeads, mHeadSize}, rt::DeviceType::kGPU, qInputDesc.type); rt::Tensor kInputTensor(const_cast(inputs[kIN_K_IDX]), - rt::Coords{runtimeBatchSize, kvSeqLen, mNumKVHeads, mHeadSize}, rt::DeviceType::kGPU, kInputDesc.type); + rt::Coords{runtimeBatchSize, runtimeSeqLen, mNumKVHeads, mHeadSize}, rt::DeviceType::kGPU, kInputDesc.type); rt::Tensor vInputTensor(const_cast(inputs[kIN_V_IDX]), - rt::Coords{runtimeBatchSize, kvSeqLen, mNumKVHeads, mHeadSize}, rt::DeviceType::kGPU, vInputDesc.type); + rt::Coords{runtimeBatchSize, runtimeSeqLen, mNumKVHeads, mHeadSize}, rt::DeviceType::kGPU, vInputDesc.type); PluginTensorDesc const& contextLengthInputDesc = inputDesc[kIN_CONTEXT_LENGTH_IDX]; rt::Tensor const contextLengthTensor(const_cast(inputs[kIN_CONTEXT_LENGTH_IDX]), @@ -854,21 +750,22 @@ int32_t AttentionPlugin::enqueue(PluginTensorDesc const* inputDesc, [[maybe_unus return 1; } - auto* alignedWorkspacePtr = static_cast(workspace); - if (alignedWorkspacePtr == nullptr - || reinterpret_cast(alignedWorkspacePtr) % static_cast(kDEVICE_ALIGNMENT) != 0) + auto const nbInputs = kNUM_REQUIRED_INPUTS + (mEnableTreeAttention ? kNUM_TREE_ATTN_OPTIONAL_INPUTS : 0); + auto const nbOutputs = kNUM_REQUIRED_OUTPUTS; + size_t space = getWorkspaceSize(inputDesc, nbInputs, outputDesc, nbOutputs); + // Align the workspace pointer so that each tensor assigned from the workspace will align to the device alignment + // granularity. + std::byte* alignedWorkspacePtr + = static_cast(std::align(kDEVICE_ALIGNMENT, space - kDEVICE_ALIGNMENT, workspace, space)); + if (alignedWorkspacePtr == nullptr) { - LOG_ERROR("Workspace pointer is not aligned to device alignment granularity"); + LOG_ERROR("Workspace size is too small to hold all data structures with correct alignment"); return 1; } - // ==================== Prefill path ==================== - // Dispatch order: sharedKV first (early return), then own-KV. - // Within each: FFPA (headSize=512 fallback) or FMHA (standard). if (executionMode == AttentionExecutionMode::kNORMAL_PREFILL || executionMode == AttentionExecutionMode::kCHUNKED_PREFILL) { - // Allocate workspace tensors for cumulative sequence lengths. rt::Tensor cuQSeqLensTensor = assignTensorFromWorkspace(alignedWorkspacePtr, {runtimeBatchSize + 1}, DataType::kINT32); @@ -883,253 +780,124 @@ int32_t AttentionPlugin::enqueue(PluginTensorDesc const* inputDesc, [[maybe_unus kernel::calCuQCuKVSeqLensAndKVEndIdxs(contextLengthTensor, kvCacheStartIdxTensor, cuQSeqLensTensor, cuKVSeqLensTensor, kvCacheEndIdxsTensor, paddedCuKVSeqLensTensor, runtimeSeqLen, stream); - // --- Shared KV prefill: Q gets RoPE, K/V read from donor layer's cache --- - if (sharedKV) - { - kernel::launchApplyRopeQOnly(ropeCosSinTensor, kvCacheEndIdxsTensor, qInputTensor, stream); + auto const fmhaMaskType = static_cast(mContextAttentionMaskType); - // Shared-KV + FFPA (headSize=512, no FMHA cubins available). - if (!mCanImplementFMHA) - { - // FFPA is a dense causal kernel — no cu_seqlens, so chunked prefill is unsupported. - if (executionMode == AttentionExecutionMode::kCHUNKED_PREFILL) - { - LOG_ERROR( - "AttentionPlugin: headSize=512 shared-KV chunked prefill is not supported with FFPA. " - "FFPA is a dense causal kernel and cannot attend to KV context longer than Q chunk length. " - "Use normal prefill (full sequence) or enable FMHA for chunked prefill."); - return -1; - } +#ifdef CUTE_DSL_FMHA_ENABLED + // Enable CuteDSL FMHA for single batch causal prefill usecase when FP8 KVCache is disabled. + // TODO: Enable multi-batch prefill and FP8 KVCache after we improve the kernel implementation. + bool const enableCuteDslFMHA = fmhaMaskType == ContextAttentionMaskType::CAUSAL && mUseCuteDslFMHA + && !mEnableFp8KVCache && runtimeBatchSize == 1; + if (enableCuteDslFMHA) + { + float const qScale = mQkvScales[0]; + int32_t const slidingWindow = mSlidingWindowSize > 0 ? mSlidingWindowSize : INT_MAX; -#ifdef CUTE_DSL_FFPA_ENABLED - if (!mCanImplementFFPA) - { - LOG_ERROR("AttentionPlugin: FFPA required for headSize=512 prefill but module failed to load."); - return -1; - } + CuteDslFMHARunner runner( + mNumQHeads, mNumKVHeads, mHeadSize, runtimeBatchSize, runtimeSeqLen, kvCacheCapacity); - LOG_DEBUG( - "AttentionPlugin: headSize=512 shared-KV prefill via FFPA native GQA " - "(B=%d, S=%d, Hq=%d, Hkv=%d, D=%d, cap=%d)", - runtimeBatchSize, runtimeSeqLen, mNumQHeads, mNumKVHeads, mHeadSize, kvCacheCapacity); + if (mEnableFp8KVCache) + { + // FP8 Q workspace: RoPE kernel quantizes roped Q to FP8 using calibrated qScale. + rt::Tensor fp8QTensor = assignTensorFromWorkspace( + alignedWorkspacePtr, {runtimeBatchSize, runtimeSeqLen, mNumQHeads, mHeadSize}, DataType::kFP8); - // Zero attention output at padding positions before FFPA writes. - // FFPA has no cu_seqlens — it processes all positions uniformly. - // BS=1 has no padding, so skip the memset. - if (runtimeBatchSize > 1) - { - zeroPrefillOutputForPaddingForFFPA( - attentionOutputTensor, runtimeBatchSize, runtimeSeqLen, mNumQHeads, mHeadSize, stream); - } + // Single kernel: RoPE Q → FP8 output, RoPE K + write FP8 K/V to cache. + kernel::launchApplyRopeWriteKVSplitQKV(ropeCosSinTensor, kvCacheEndIdxsTensor, qInputTensor, + kInputTensor, vInputTensor, kvCacheTensor, kScale, vScale, stream, fp8QTensor.rawPointer(), qScale); - // Extract K/V from donor's cache (compact: only first runtimeSeqLen tokens) - // and run FFPA with native GQA. seqlenK = runtimeSeqLen so that FFPA's - // bottom-right causal mask offset (seqlenK - seqlenQ) is 0, correctly - // bounding attention to valid positions. The compact deinterleave produces - // [B, runtimeSeqLen, Hkv, D] so physical stride matches seqlenK — no batch - // stride override needed. - auto [kSplit, vSplit] = deinterleaveKVCache(kvCacheTensor, alignedWorkspacePtr, runtimeBatchSize, - mNumKVHeads, kvCacheCapacity, mHeadSize, runtimeSeqLen, stream); - dispatchFFPAKernel(qInputTensor.dataPointer(), kSplit.dataPointer(), - vSplit.dataPointer(), attentionOutputTensor.dataPointer(), runtimeBatchSize, - runtimeSeqLen, runtimeSeqLen, mNumQHeads, mNumKVHeads, mHeadSize, stream); -#else - LOG_ERROR("AttentionPlugin: headSize=512 shared-KV prefill requires FFPA (CUTE_DSL_FFPA_ENABLED)."); - return -1; -#endif - return 0; + runner.run(fp8QTensor.rawPointer(), // Q [b, s_q, h_q, d] FP8 + kvCacheTensor.rawPointer(), // KV [b, 2, h_k, cap, d] FP8 + attentionOutputTensor.dataPointer(), // O [b, s_q, h_q, d] FP16 + paddedCuKVSeqLensTensor.dataPointer(), // cu_kv_seqlens [b+1] + stream, slidingWindow, /*fp8Input=*/true, qScale, kScale, vScale); } - - // Shared-KV + FMHA: read Q/KV directly from donor's cache. -#ifdef CUTE_DSL_FMHA_ENABLED - if (mUseCuteDslFMHA) + else { - // CuTe DSL FMHA reads interleaved KV cache natively. - int32_t const slidingWindow = mSlidingWindowSize > 0 ? mSlidingWindowSize : INT_MAX; - CuteDslFMHARunner runner( - mNumQHeads, mNumKVHeads, mHeadSize, runtimeBatchSize, runtimeSeqLen, kvCacheCapacity); + // FP16 path: RoPE Q in-place, write FP16 K/V to cache. + kernel::launchApplyRopeWriteKVSplitQKV(ropeCosSinTensor, kvCacheEndIdxsTensor, qInputTensor, + kInputTensor, vInputTensor, kvCacheTensor, kScale, vScale, stream); + runner.run(qInputTensor.dataPointer(), // Q [b, s_q, h_q, d] - kvCacheTensor.dataPointer(), // KV [b, 2, h_k, cap, d] (donor's cache) + kvCacheTensor.dataPointer(), // KV [b, 2, h_k, cap, d] attentionOutputTensor.dataPointer(), // O [b, s_q, h_q, d] paddedCuKVSeqLensTensor.dataPointer(), // cu_kv_seqlens [b+1] stream, slidingWindow); } - else -#endif - { - // FMHA_v2 requires separate K/V — deinterleave from donor's cache. - auto fmhaRunner = ContextFMHARunner(mDataType, runtimeBatchSize, runtimeSeqLen, mNumQHeads, mNumKVHeads, - mHeadSize, mSMVersion, AttentionInputLayout::SEPARATE_Q_K_V); - FusedMultiheadAttentionParamsV2 params{}; - fmhaRunner.setupParams(params); - params.cu_q_seqlens = cuQSeqLensTensor.dataPointer(); - - // Normal prefill: compact deinterleave (seqLen tokens) so FMHA_v2's - // s_kv-derived batch stride matches the physical layout. - // Chunked prefill: full deinterleave, s_kv = kvCacheCapacity. - bool const compact = (executionMode == AttentionExecutionMode::kNORMAL_PREFILL); - int32_t const seqLen = compact ? runtimeSeqLen : 0; - auto [kSplit, vSplit] = deinterleaveKVCache(kvCacheTensor, alignedWorkspacePtr, runtimeBatchSize, - mNumKVHeads, kvCacheCapacity, mHeadSize, seqLen, stream); - - params.s_kv = compact ? runtimeSeqLen : kvCacheCapacity; - params.q_ptr = qInputTensor.dataPointer(); - params.k_ptr = kSplit.dataPointer(); - params.v_ptr = vSplit.dataPointer(); - params.cu_kv_seqlens = cuKVSeqLensTensor.dataPointer(); - params.o_ptr = attentionOutputTensor.dataPointer(); - fmhaRunner.dispatchFMHAKernel(params, stream); - } - return 0; } - - // --- Own KV prefill: RoPE Q+K, write K/V to cache, then run attention kernel --- - - // Own-KV + FFPA (headSize=512, no FMHA cubins available). - if (!mCanImplementFMHA) + else +#endif { - // RoPE + write KV to cache. writeKInPlace=true so kInput gets roped for FFPA below. - kernel::launchApplyRopeWriteKV(ropeCosSinTensor, kvCacheEndIdxsTensor, qInputTensor, kInputTensor, - vInputTensor, kvCacheTensor, kScale, vScale, stream, true); + auto fmhaRunner = ContextFMHARunner(mDataType, runtimeBatchSize, runtimeSeqLen, mNumQHeads, mNumKVHeads, + mHeadSize, mSMVersion, AttentionInputLayout::SEPARATE_Q_K_V, fmhaMaskType); -#ifdef CUTE_DSL_FFPA_ENABLED - if (!mCanImplementFFPA) - { - LOG_ERROR("AttentionPlugin: FFPA required for headSize=512 prefill but module failed to load."); - return -1; - } + // Prepare FMHA_v2 params to launch FMHA kernel + FusedMultiheadAttentionParamsV2 params{}; + fmhaRunner.setupParams(params); + params.cu_q_seqlens = cuQSeqLensTensor.dataPointer(); - // Use FFPA d512 causal kernel with native GQA support (no K/V expansion needed). - LOG_DEBUG( - "AttentionPlugin: headSize=512 own-KV prefill via FFPA native GQA " - "(B=%d, S=%d, Hq=%d, Hkv=%d, D=%d)", - runtimeBatchSize, runtimeSeqLen, mNumQHeads, mNumKVHeads, mHeadSize); - - dispatchFFPAKernel(qInputTensor.dataPointer(), kInputTensor.dataPointer(), - vInputTensor.dataPointer(), attentionOutputTensor.dataPointer(), runtimeBatchSize, - runtimeSeqLen, runtimeSeqLen, mNumQHeads, mNumKVHeads, mHeadSize, stream); -#else - LOG_ERROR("AttentionPlugin: headSize=512 own-KV prefill requires FFPA (CUTE_DSL_FFPA_ENABLED)."); - return -1; -#endif - } - // Own-KV + FMHA (standard path). - else - { -#ifdef CUTE_DSL_FMHA_ENABLED - if (mUseCuteDslFMHA) + if (executionMode == AttentionExecutionMode::kCHUNKED_PREFILL) { - // CuTe DSL FMHA uses SplitQKV RoPE variant that writes K/V to interleaved cache. - float const qScale = mQkvScales[0]; - int32_t const slidingWindow = mSlidingWindowSize > 0 ? mSlidingWindowSize : INT_MAX; - - CuteDslFMHARunner runner( - mNumQHeads, mNumKVHeads, mHeadSize, runtimeBatchSize, runtimeSeqLen, kvCacheCapacity); + // kvCache: [b, 2, hkv, s, d] -> split K [b, s, hkv, d] + V [b, s, hkv, d] + kernel::launchApplyRopeWriteKV(ropeCosSinTensor, kvCacheEndIdxsTensor, qInputTensor, kInputTensor, + vInputTensor, kvCacheTensor, kScale, vScale, stream, false); - if (mEnableFp8KVCache) - { - // FP8: RoPE quantizes Q→FP8, writes FP8 K/V to cache. - rt::Tensor fp8QTensor = assignTensorFromWorkspace( - alignedWorkspacePtr, {runtimeBatchSize, runtimeSeqLen, mNumQHeads, mHeadSize}, DataType::kFP8); - - // Single kernel: RoPE Q → FP8 output, RoPE K + write FP8 K/V to cache. - kernel::launchApplyRopeWriteKVSplitQKV(ropeCosSinTensor, kvCacheEndIdxsTensor, qInputTensor, - kInputTensor, vInputTensor, kvCacheTensor, kScale, vScale, stream, fp8QTensor.rawPointer(), - qScale); - - runner.run(fp8QTensor.rawPointer(), // Q [b, s_q, h_q, d] FP8 - kvCacheTensor.rawPointer(), // KV [b, 2, h_k, cap, d] FP8 - attentionOutputTensor.dataPointer(), // O [b, s_q, h_q, d] FP16 - paddedCuKVSeqLensTensor.dataPointer(), // cu_kv_seqlens [b+1] - stream, slidingWindow, /*fp8Input=*/true, qScale, kScale, vScale); - } - else - { - // FP16 path: RoPE Q in-place, write FP16 K/V to cache. - kernel::launchApplyRopeWriteKVSplitQKV(ropeCosSinTensor, kvCacheEndIdxsTensor, qInputTensor, - kInputTensor, vInputTensor, kvCacheTensor, kScale, vScale, stream); - - runner.run(qInputTensor.dataPointer(), // Q [b, s_q, h_q, d] - kvCacheTensor.dataPointer(), // KV [b, 2, h_k, cap, d] - attentionOutputTensor.dataPointer(), // O [b, s_q, h_q, d] - paddedCuKVSeqLensTensor.dataPointer(), // cu_kv_seqlens [b+1] - stream, slidingWindow); - } + // Allocate a single workspace and split into K and V halves by pointer arithmetic. + rt::Tensor kvWorkspaceTensor = assignTensorFromWorkspace(alignedWorkspacePtr, + {runtimeBatchSize, 2, mNumKVHeads, kvCacheCapacity, mHeadSize}, DataType::kHALF); + size_t const halfSize + = static_cast(runtimeBatchSize) * kvCacheCapacity * mNumKVHeads * mHeadSize; + half* kvWorkspacePtr = kvWorkspaceTensor.dataPointer(); + rt::Tensor kWorkspaceTensor(kvWorkspacePtr, + rt::Coords{runtimeBatchSize, kvCacheCapacity, mNumKVHeads, mHeadSize}, rt::DeviceType::kGPU, + DataType::kHALF); + rt::Tensor vWorkspaceTensor(kvWorkspacePtr + halfSize, + rt::Coords{runtimeBatchSize, kvCacheCapacity, mNumKVHeads, mHeadSize}, rt::DeviceType::kGPU, + DataType::kHALF); + // seqLen == 0: copy the full KV-cache capacity (matches the pre-seqLen kernel behavior). + kernel::cvtKVLayoutBHSDToSplitKV( + kvCacheTensor, kWorkspaceTensor, vWorkspaceTensor, rt::Tensor{}, /*seqLen=*/0, stream); + + // Set device ptr for FMHA kernel. + params.s_kv = kvCacheCapacity; + params.q_ptr = qInputTensor.dataPointer(); + params.k_ptr = kWorkspaceTensor.dataPointer(); + params.v_ptr = vWorkspaceTensor.dataPointer(); + params.cu_kv_seqlens = cuKVSeqLensTensor.dataPointer(); + params.o_ptr = attentionOutputTensor.dataPointer(); } else -#endif - { - // FMHA_v2 fallback: separate K/V pointers required. - auto fmhaRunner = ContextFMHARunner(mDataType, runtimeBatchSize, runtimeSeqLen, mNumQHeads, mNumKVHeads, - mHeadSize, mSMVersion, AttentionInputLayout::SEPARATE_Q_K_V); - FusedMultiheadAttentionParamsV2 params{}; - fmhaRunner.setupParams(params); - params.cu_q_seqlens = cuQSeqLensTensor.dataPointer(); - - if (executionMode == AttentionExecutionMode::kCHUNKED_PREFILL) - { - // Chunked: RoPE + write to cache, then deinterleave for FMHA_v2 input. - kernel::launchApplyRopeWriteKV(ropeCosSinTensor, kvCacheEndIdxsTensor, qInputTensor, kInputTensor, - vInputTensor, kvCacheTensor, kScale, vScale, stream, false); - - auto [kSplit, vSplit] = deinterleaveKVCache(kvCacheTensor, alignedWorkspacePtr, runtimeBatchSize, - mNumKVHeads, kvCacheCapacity, mHeadSize, 0, stream); - - params.s_kv = kvCacheCapacity; - params.q_ptr = qInputTensor.dataPointer(); - params.k_ptr = kSplit.dataPointer(); - params.v_ptr = vSplit.dataPointer(); - params.cu_kv_seqlens = cuKVSeqLensTensor.dataPointer(); - params.o_ptr = attentionOutputTensor.dataPointer(); - } - else - { - // Normal prefill: RoPE in-place, read K/V directly from input tensors. - kernel::launchApplyRopeWriteKV(ropeCosSinTensor, std::nullopt, qInputTensor, kInputTensor, - vInputTensor, kvCacheTensor, kScale, vScale, stream, true); - - params.s_kv = runtimeSeqLen; - params.q_ptr = qInputTensor.dataPointer(); - params.k_ptr = kInputTensor.dataPointer(); - params.v_ptr = vInputTensor.dataPointer(); - params.cu_kv_seqlens = cuQSeqLensTensor.dataPointer(); - params.o_ptr = attentionOutputTensor.dataPointer(); - } + { // SEPARATE_Q_K_V + kernel::launchApplyRopeWriteKV(ropeCosSinTensor, std::nullopt, qInputTensor, kInputTensor, vInputTensor, + kvCacheTensor, kScale, vScale, stream, true); - fmhaRunner.dispatchFMHAKernel(params, stream); + params.s_kv = runtimeSeqLen; + params.q_ptr = qInputTensor.dataPointer(); + params.k_ptr = kInputTensor.dataPointer(); + params.v_ptr = vInputTensor.dataPointer(); + params.cu_kv_seqlens = cuQSeqLensTensor.dataPointer(); + params.o_ptr = attentionOutputTensor.dataPointer(); } + + // Dispatch FMHA kernel + fmhaRunner.dispatchFMHAKernel(params, stream); } } - // ==================== Decode path (vanilla or tree) ==================== else { - // RoPE setup: sharedKV → Q only, own-KV → Q+K with KV cache write. + // Prepare Decoding attention runner parameter to dispatch kernel if (executionMode == AttentionExecutionMode::kTREE_DECODING) { - if (sharedKV) - { - kernel::launchApplyRopeQOnlyTreeDecoding(ropeCosSinTensor, attentionPosIdTensor, qInputTensor, stream); - } - else - { - kernel::launchApplyRopeWriteKVTreeDecoding(ropeCosSinTensor, contextLengthTensor, attentionPosIdTensor, - qInputTensor, kInputTensor, vInputTensor, kvCacheTensor, kScale, vScale, stream); - } + // Execute tree attention decoding. + kernel::launchApplyRopeWriteKVTreeDecoding(ropeCosSinTensor, contextLengthTensor, attentionPosIdTensor, + qInputTensor, kInputTensor, vInputTensor, kvCacheTensor, kScale, vScale, stream); } else { - if (sharedKV) - { - kernel::launchApplyRopeQOnly(ropeCosSinTensor, contextLengthTensor, qInputTensor, stream); - } - else - { - kernel::launchApplyRopeWriteKV(ropeCosSinTensor, contextLengthTensor, qInputTensor, kInputTensor, - vInputTensor, kvCacheTensor, kScale, vScale, stream, false); - } + // Execute vanilla decoding. + kernel::launchApplyRopeWriteKV(ropeCosSinTensor, contextLengthTensor, qInputTensor, kInputTensor, + vInputTensor, kvCacheTensor, kScale, vScale, stream, false); } - // XQA decode kernel dispatch. auto xqaRunner = DecoderXQARunner(mDataType, selectKvCacheDataType(mEnableFp8KVCache), runtimeBatchSize, mNumQHeads, mNumKVHeads, mHeadSize, mSMVersion); XQALaunchParams params = xqaRunner.initXQAParams(); @@ -1143,7 +911,6 @@ int32_t AttentionPlugin::enqueue(PluginTensorDesc const* inputDesc, [[maybe_unus params.kvCache.data = kvCacheTensor.rawPointer(); params.kvCache.sequence_lengths = contextLengthTensor.dataPointer(); params.kvCache.capacity = kvCacheCapacity; - params.slidingWinSize = mSlidingWindowSize > 0 ? static_cast(mSlidingWindowSize) : 0U; if (executionMode == AttentionExecutionMode::kTREE_DECODING) { // Execute tree attention decoding. @@ -1160,36 +927,44 @@ int32_t AttentionPlugin::enqueue(PluginTensorDesc const* inputDesc, [[maybe_unus return 0; } -int32_t AttentionPlugin::onShapeChange([[maybe_unused]] PluginTensorDesc const* in, [[maybe_unused]] int32_t nbInputs, - [[maybe_unused]] PluginTensorDesc const* out, [[maybe_unused]] int32_t nbOutputs) noexcept +size_t AttentionPlugin::getSerializationSize() const noexcept { - return 0; + return sizeof(mNumQHeads) + sizeof(mNumKVHeads) + sizeof(mHeadSize) + sizeof(mEnableTreeAttention) + + sizeof(mEnableFp8KVCache) + sizeof(int32_t) /* slidingWindowSize */ + + sizeof(int32_t) /* qkv scale count */ + mQkvScales.size() * sizeof(float) + + sizeof(int32_t) /* context mask serialization tag */ + sizeof(mContextAttentionMaskType); } -IPluginV3* AttentionPlugin::attachToContext([[maybe_unused]] IPluginResourceContext* context) noexcept +void AttentionPlugin::serialize(void* buffer) const noexcept { - return clone(); + std::byte* byteBuffer = static_cast(buffer); + serializeValue(&byteBuffer, mNumQHeads); + serializeValue(&byteBuffer, mNumKVHeads); + serializeValue(&byteBuffer, mHeadSize); + serializeValue(&byteBuffer, mEnableTreeAttention); + serializeValue(&byteBuffer, mEnableFp8KVCache); + serializeValue(&byteBuffer, mSlidingWindowSize); + int32_t const qkvScaleCount = static_cast(mQkvScales.size()); + serializeValue(&byteBuffer, qkvScaleCount); + for (auto const& s : mQkvScales) + { + serializeValue(&byteBuffer, s); + } + serializeValue(&byteBuffer, kContextMaskSerializationTag); + serializeValue(&byteBuffer, mContextAttentionMaskType); } -PluginFieldCollection const* AttentionPlugin::getFieldsToSerialize() noexcept +int32_t AttentionPlugin::initialize() noexcept { - mDataToSerialize.clear(); - mDataToSerialize.emplace_back("num_q_heads", &mNumQHeads, PluginFieldType::kINT32, 1); - mDataToSerialize.emplace_back("num_kv_heads", &mNumKVHeads, PluginFieldType::kINT32, 1); - mDataToSerialize.emplace_back("head_size", &mHeadSize, PluginFieldType::kINT32, 1); - mDataToSerialize.emplace_back("enable_tree_attention", &mEnableTreeAttention, PluginFieldType::kINT32, 1); - mDataToSerialize.emplace_back("enable_fp8_kv_cache", &mEnableFp8KVCache, PluginFieldType::kINT32, 1); - mDataToSerialize.emplace_back("sliding_window_size", &mSlidingWindowSize, PluginFieldType::kINT32, 1); - mDataToSerialize.emplace_back( - "qkv_scales", mQkvScales.data(), PluginFieldType::kFLOAT32, static_cast(mQkvScales.size())); - mFCToSerialize.nbFields = static_cast(mDataToSerialize.size()); - mFCToSerialize.fields = mDataToSerialize.data(); - return &mFCToSerialize; + return 0; } -// --------------------------------------------------------------------------- -// Creator -// --------------------------------------------------------------------------- +void AttentionPlugin::terminate() noexcept {} + +void AttentionPlugin::destroy() noexcept +{ + delete this; +} AttentionPluginCreator::AttentionPluginCreator() { @@ -1205,6 +980,9 @@ AttentionPluginCreator::AttentionPluginCreator() mPluginAttributes.emplace_back(PluginField("enable_fp8_kv_cache", nullptr, PluginFieldType::kINT32, 0)); // Sliding window size (-1 = no sliding window, >0 = window size) mPluginAttributes.emplace_back(PluginField("sliding_window_size", nullptr, PluginFieldType::kINT32, 0)); + // Optional context attention mask type (ContextAttentionMaskType: 0 = PADDING, 1 = CAUSAL, + // 2 = SLIDING_OR_CHUNKED_CAUSAL, 3 = CUSTOM_MASK). Defaults to CAUSAL when unset. + mPluginAttributes.emplace_back(PluginField("context_attention_mask_type", nullptr, PluginFieldType::kINT32, 0)); // Optional QKV dequant scales [q, k, v] for FP8 attention mPluginAttributes.emplace_back(PluginField("qkv_scales", nullptr, PluginFieldType::kFLOAT32, 0)); // Enforce Core parameters are specified. @@ -1217,14 +995,14 @@ char const* AttentionPluginCreator::getPluginName() const noexcept return kATTENTION_PLUGIN_NAME; } -PluginFieldCollection const* AttentionPluginCreator::getFieldNames() noexcept +nvinfer1::PluginFieldCollection const* AttentionPluginCreator::getFieldNames() noexcept { return &mFieldCollection; } void AttentionPluginCreator::setPluginNamespace(char const* libNamespace) noexcept { - mNamespace = libNamespace ? libNamespace : ""; + mNamespace = libNamespace; } char const* AttentionPluginCreator::getPluginNamespace() const noexcept @@ -1237,13 +1015,70 @@ char const* AttentionPluginCreator::getPluginVersion() const noexcept return kATTENTION_PLUGIN_VERSION; } -IPluginV3* AttentionPluginCreator::createPlugin( - char const* name, PluginFieldCollection const* fc, [[maybe_unused]] TensorRTPhase phase) noexcept +nvinfer1::IPluginV2* AttentionPluginCreator::createPlugin( + char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept { try { - auto* plugin = new AttentionPlugin(std::string(name), fc); - plugin->setPluginNamespace(mNamespace.c_str()); + std::optional numQHeads = parsePluginScalarField("num_q_heads", fc); + std::optional numKVHeads = parsePluginScalarField("num_kv_heads", fc); + std::optional headSize = parsePluginScalarField("head_size", fc); + std::optional enableTreeAttention = parsePluginScalarField("enable_tree_attention", fc); + std::optional enableFp8KVCache = parsePluginScalarField("enable_fp8_kv_cache", fc); + // sliding_window_size: -1 = no sliding window (default), >0 = sliding window size + int32_t slidingWindowSize = parsePluginScalarField("sliding_window_size", fc).value_or(-1); + // Make enable_fp8_kv_cache optional with default value 0 (disable by default) + int32_t enableFp8KVCacheValue = enableFp8KVCache.value_or(0); + int32_t contextAttentionMaskType = parsePluginScalarField("context_attention_mask_type", fc) + .value_or(static_cast(ContextAttentionMaskType::CAUSAL)); + if (!parsePluginScalarField("context_attention_mask_type", fc).has_value()) + { + if (auto legacyBidirectionalPrefill = parsePluginScalarField("enable_bidirectional_prefill", fc)) + { + contextAttentionMaskType = (legacyBidirectionalPrefill.value() != 0) + ? static_cast(ContextAttentionMaskType::PADDING) + : static_cast(ContextAttentionMaskType::CAUSAL); + } + } + if (contextAttentionMaskType < static_cast(ContextAttentionMaskType::PADDING) + || contextAttentionMaskType > static_cast(ContextAttentionMaskType::CUSTOM_MASK)) + { + LOG_ERROR("Invalid context_attention_mask_type %d (expected 0-3).", contextAttentionMaskType); + return nullptr; + } + + bool checkRequiredFields = numQHeads.has_value() && headSize.has_value() && numKVHeads.has_value() + && enableTreeAttention.has_value(); + if (!checkRequiredFields) + { + LOG_ERROR("Missing required AttentionPlugin fields."); + return nullptr; + } + + std::vector qkvScales; + for (int32_t i = 0; i < fc->nbFields; ++i) + { + if (std::string("qkv_scales") == fc->fields[i].name) + { + auto const* data = static_cast(fc->fields[i].data); + qkvScales.assign(data, data + fc->fields[i].length); + break; + } + } + + if (enableFp8KVCacheValue && qkvScales.size() != 3) + { + LOG_ERROR( + "FP8 KV cache enabled but qkv_scales has %zu elements (expected 3). " + "Re-export the model to include QKV scales [q, k, v].", + qkvScales.size()); + return nullptr; + } + + AttentionPlugin* plugin = new AttentionPlugin(std::string(name), numQHeads.value(), numKVHeads.value(), + headSize.value(), enableTreeAttention.value(), enableFp8KVCacheValue, slidingWindowSize, qkvScales, + contextAttentionMaskType); + return plugin; } catch (std::exception const& e) @@ -1253,5 +1088,19 @@ IPluginV3* AttentionPluginCreator::createPlugin( return nullptr; } +nvinfer1::IPluginV2* AttentionPluginCreator::deserializePlugin( + char const* name, void const* serialData, size_t serialLength) noexcept +{ + try + { + return new AttentionPlugin(name, static_cast(serialData), serialLength); + } + catch (std::exception const& e) + { + LOG_ERROR("Failed to deserialize AttentionPlugin: %s", e.what()); + } + return nullptr; +} + } // namespace plugins } // namespace trt_edgellm diff --git a/cpp/plugins/attentionPlugin/attentionPlugin.h b/cpp/plugins/attentionPlugin/attentionPlugin.h index 4c81c9a0..fd54406a 100644 --- a/cpp/plugins/attentionPlugin/attentionPlugin.h +++ b/cpp/plugins/attentionPlugin/attentionPlugin.h @@ -21,24 +21,18 @@ #include #include #include -#include #include -#include "common/tensor.h" - namespace trt_edgellm { namespace plugins { -//! \brief TensorRT plugin for attention operations (V3 — IPluginV3). +//! \brief TensorRT plugin for attention operations (context and decode) //! //! This plugin implements efficient attention mechanisms including context attention (prefill) //! and decode attention with KV cache support. -class AttentionPlugin : public nvinfer1::IPluginV3, - public nvinfer1::IPluginV3OneCore, - public nvinfer1::IPluginV3OneBuildV2, - public nvinfer1::IPluginV3OneRuntime +class AttentionPlugin : public nvinfer1::IPluginV2DynamicExt { public: //! \brief Constructor for attention plugin with configuration parameters @@ -50,67 +44,125 @@ class AttentionPlugin : public nvinfer1::IPluginV3, //! \param[in] enableFp8KVCache Whether to enable FP8 KV cache //! \param[in] slidingWindowSize Sliding window size (-1 = no sliding window) //! \param[in] qkvScales Optional [q, k, v] FP8 dequant scales (required when enableFp8KVCache) + //! \param[in] contextAttentionMaskType Context prefill mask type (ContextAttentionMaskType enum value: + //! 0 = PADDING, 1 = CAUSAL, 2 = SLIDING_OR_CHUNKED_CAUSAL, 3 = CUSTOM_MASK). Defaults to CAUSAL. AttentionPlugin(std::string const& name, int32_t numQHeads, int32_t numKVHeads, int32_t headSize, int32_t supportsSpecDecode, int32_t enableFp8KVCache, int32_t slidingWindowSize = -1, - std::vector const& qkvScales = {}); - AttentionPlugin(std::string const& name, nvinfer1::PluginFieldCollection const* fc); + std::vector const& qkvScales = {}, int32_t contextAttentionMaskType = 1); + + //! \brief Constructor for deserialization + //! \param[in] name Plugin instance name + //! \param[in] data Serialized plugin data + //! \param[in] length Length of serialized data + AttentionPlugin(std::string const& name, std::byte const* data, size_t length); + //! Force to distinguish different instances of the plugin AttentionPlugin() = delete; + AttentionPlugin(AttentionPlugin const&) = delete; + ~AttentionPlugin() override; - // IPluginV3 - nvinfer1::IPluginCapability* getCapabilityInterface(nvinfer1::PluginCapabilityType type) noexcept override; - nvinfer1::IPluginV3* clone() noexcept override; + //! \name IPluginV2DynamicExt Methods + //! @{ - // IPluginV3OneCore - char const* getPluginName() const noexcept override; - char const* getPluginVersion() const noexcept override; - char const* getPluginNamespace() const noexcept override; + //! \brief Clone the plugin instance + //! \return Pointer to cloned plugin + nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - // IPluginV3OneBuild + //! \brief Get number of outputs + //! \return Number of output tensors int32_t getNbOutputs() const noexcept override; - int32_t getOutputDataTypes(nvinfer1::DataType* outputTypes, int32_t nbOutputs, nvinfer1::DataType const* inputTypes, - int32_t nbInputs) const noexcept override; - int32_t getOutputShapes(nvinfer1::DimsExprs const* inputs, int32_t nbInputs, nvinfer1::DimsExprs const* shapeInputs, - int32_t nbShapeInputs, nvinfer1::DimsExprs* outputs, int32_t nbOutputs, + + //! \brief Get output data type + //! \param[in] index Output index + //! \param[in] inputTypes Array of input data types + //! \param[in] nbInputs Number of inputs + //! \return Output data type + nvinfer1::DataType getOutputDataType( + int32_t index, nvinfer1::DataType const* inputTypes, int32_t nbInputs) const noexcept override; + + //! \brief Get output dimensions + //! \param[in] outputIndex Output tensor index + //! \param[in] inputs Input tensor dimensions + //! \param[in] nbInputs Number of inputs + //! \param[in] exprBuilder Expression builder for dimension calculations + //! \return Output tensor dimensions + nvinfer1::DimsExprs getOutputDimensions(int32_t outputIndex, nvinfer1::DimsExprs const* inputs, int32_t nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination(int32_t pos, nvinfer1::DynamicPluginTensorDesc const* inOut, int32_t nbInputs, - int32_t nbOutputs) noexcept override; - int32_t configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int32_t nbInputs, + + //! \brief Check if format combination is supported + //! \param[in] pos Position in the input/output tensor list + //! \param[in] inOut Array of input and output tensor descriptors + //! \param[in] nbInputs Number of inputs + //! \param[in] nbOutputs Number of outputs + //! \return True if format combination is supported + bool supportsFormatCombination( + int32_t pos, nvinfer1::PluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept override; + + //! \brief Configure the plugin with input and output tensors + //! \param[in] in Input tensor descriptors + //! \param[in] nbInputs Number of inputs + //! \param[in] out Output tensor descriptors + //! \param[in] nbOutputs Number of outputs + void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int32_t nbInputs, nvinfer1::DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::DynamicPluginTensorDesc const* inputs, int32_t nbInputs, - nvinfer1::DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept override; - int32_t getAliasedInput(int32_t outputIndex) noexcept override; - // IPluginV3OneRuntime + //! \brief Get workspace size required by the plugin + //! \param[in] inputs Input tensor descriptors + //! \param[in] nbInputs Number of inputs + //! \param[in] outputs Output tensor descriptors + //! \param[in] nbOutputs Number of outputs + //! \return Workspace size in bytes + size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int32_t nbInputs, + nvinfer1::PluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept override; + + //! \brief Execute the plugin + //! \param[in] inputDesc Input tensor descriptors + //! \param[in] outputDesc Output tensor descriptors + //! \param[in] inputs Input tensor data pointers + //! \param[out] outputs Output tensor data pointers + //! \param[in] workspace Workspace memory pointer + //! \param[in] stream CUDA stream for execution + //! \return 0 on success, non-zero on failure int32_t enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - int32_t onShapeChange(nvinfer1::PluginTensorDesc const* in, int32_t nbInputs, nvinfer1::PluginTensorDesc const* out, - int32_t nbOutputs) noexcept override; - nvinfer1::IPluginV3* attachToContext(nvinfer1::IPluginResourceContext* context) noexcept override; - nvinfer1::PluginFieldCollection const* getFieldsToSerialize() noexcept override; + //! \brief Get serialization size + //! \return Size in bytes required for serialization + size_t getSerializationSize() const noexcept override; + + //! \brief Serialize the plugin + //! \param[out] buffer Buffer to write serialized data + void serialize(void* buffer) const noexcept override; + + //! \brief Get plugin type + //! \return Plugin type string + char const* getPluginType() const noexcept override; + + //! \brief Get plugin namespace + //! \return Plugin namespace string + char const* getPluginNamespace() const noexcept override; + + //! \brief Set plugin namespace + //! \param[in] pluginNamespace Namespace to set void setPluginNamespace(char const* pluginNamespace) noexcept; -private: - //! Split a BHSD-layout KV cache [B, 2, Hkv, cap, D] into separate K and V tensors. - //! When seqLen == 0 (default), copies the full capacity → output is [B, cap, Hkv, D]. - //! When seqLen > 0, copies only the first seqLen tokens → output is [B, seqLen, Hkv, D]. - //! The compact form allows downstream kernels to derive batch stride from the output's S dimension. - static std::pair deinterleaveKVCache(rt::Tensor const& kvCacheTensor, - std::byte*& workspacePtr, int32_t batchSize, int32_t numKVHeads, int32_t kvCacheCapacity, int32_t headSize, - int32_t seqLen, cudaStream_t stream); - - //! Launch the CuTe DSL FFPA d512 causal attention kernel. - static void dispatchFFPAKernel(half const* q, half const* k, half const* v, half* o, int32_t batchSize, - int32_t seqlenQ, int32_t seqlenK, int32_t numQHeads, int32_t numKVHeads, int32_t headDim, cudaStream_t stream); - - //! Zero the attention output buffer before FFPA prefill. - //! FFPA is a dense causal kernel with no cu_seqlens support, so it processes padding positions as real data. - //! Zeroing the output ensures padding positions don't carry NaN/garbage into downstream layers. - static void zeroPrefillOutputForPaddingForFFPA(rt::Tensor& attentionOutput, int32_t batchSize, int32_t seqLen, - int32_t numQHeads, int32_t headSize, cudaStream_t stream); + //! \brief Get plugin version + //! \return Plugin version string + char const* getPluginVersion() const noexcept override; + + //! \brief Initialize the plugin + //! \return 0 on success, non-zero on failure + int32_t initialize() noexcept override; + + //! \brief Terminate the plugin and release resources + void terminate() noexcept override; + + //! \brief Destroy the plugin instance + void destroy() noexcept override; + + //! @} protected: std::string mLayerName; //!< Plugin layer name @@ -140,45 +192,63 @@ class AttentionPlugin : public nvinfer1::IPluginV3, //! Sliding window size for attention (-1 = no sliding window, >0 = window size) int32_t mSlidingWindowSize = -1; + //! Context attention mask type for prefill, stored as a ContextAttentionMaskType enum value. + //! 0 = PADDING (bidirectional full-prefix), 1 = CAUSAL (default), 2 = SLIDING_OR_CHUNKED_CAUSAL, 3 = CUSTOM_MASK. + int32_t mContextAttentionMaskType{1}; + #ifdef CUTE_DSL_FMHA_ENABLED bool mUseCuteDslFMHA{true}; #else bool mUseCuteDslFMHA{false}; #endif - - //! Whether FMHA context kernels are available for this configuration. - //! When false (e.g. headSize=512), the prefill path uses XQA instead. - bool mCanImplementFMHA{true}; - - //! Whether FFPA d512 kernel is available for headSize=512 prefill+decode. - bool mCanImplementFFPA{false}; - - //! Whether XQA decode kernels are available. - bool mCanImplementXQA{false}; - - std::vector mDataToSerialize; - nvinfer1::PluginFieldCollection mFCToSerialize{}; }; //! \brief Factory class for creating AttentionPlugin instances -class AttentionPluginCreator : public nvinfer1::IPluginCreatorV3One +class AttentionPluginCreator : public nvinfer1::IPluginCreator { public: AttentionPluginCreator(); + ~AttentionPluginCreator() override = default; + //! \brief Get plugin name + //! \return Plugin name string char const* getPluginName() const noexcept override; - char const* getPluginVersion() const noexcept override; + + //! \brief Get plugin field collection + //! \return Pointer to plugin field collection containing all plugin fields nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - char const* getPluginNamespace() const noexcept override; + + //! \brief Set plugin namespace + //! \param[in] pluginNamespace Namespace to set void setPluginNamespace(char const* pluginNamespace) noexcept; - nvinfer1::IPluginV3* createPlugin( - char const* name, nvinfer1::PluginFieldCollection const* fc, nvinfer1::TensorRTPhase phase) noexcept override; + + //! \brief Get plugin namespace + //! \return Plugin namespace string + char const* getPluginNamespace() const noexcept override; + + //! \brief Get plugin version + //! \return Plugin version string + char const* getPluginVersion() const noexcept override; + + //! \brief Create a new plugin instance + //! \param[in] name Plugin instance name + //! \param[in] fc Plugin field collection containing configuration parameters + //! \return Pointer to created plugin instance + nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; + + //! \brief Deserialize a plugin instance from data + //! \param[in] name Plugin instance name + //! \param[in] serialData Serialized plugin data + //! \param[in] serialLength Length of serialized data in bytes + //! \return Pointer to deserialized plugin instance + nvinfer1::IPluginV2* deserializePlugin( + char const* name, void const* serialData, size_t serialLength) noexcept override; private: - static nvinfer1::PluginFieldCollection mFieldCollection; - static std::vector mPluginAttributes; - std::string mNamespace; + static nvinfer1::PluginFieldCollection mFieldCollection; //!< Plugin field collection for registration + static std::vector mPluginAttributes; //!< Plugin attributes/fields + std::string mNamespace; //!< Plugin namespace }; } // namespace plugins diff --git a/cpp/runtime/audioDecodeRunner.cpp b/cpp/runtime/audioDecodeRunner.cpp new file mode 100644 index 00000000..fc66b6da --- /dev/null +++ b/cpp/runtime/audioDecodeRunner.cpp @@ -0,0 +1,421 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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. + */ + +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/audioDecodeRunner.h" + +#include "common/cudaUtils.h" +#include "common/logger.h" +#include "common/mmapReader.h" + +#include +#include +#include +#include + +namespace trt_edgellm +{ +namespace rt +{ +namespace +{ + +rt::Coords coordsFromJson(nlohmann::json const& shape) +{ + std::vector dims; + dims.reserve(shape.size()); + for (auto const& dim : shape) + { + dims.push_back(dim.get()); + } + return rt::Coords(dims); +} + +nvinfer1::DataType dataTypeFromTorchString(std::string const& dtype) +{ + if (dtype == "torch.float16" || dtype == "float16" || dtype == "fp16") + { + return nvinfer1::DataType::kHALF; + } + if (dtype == "torch.float32" || dtype == "float32" || dtype == "fp32") + { + return nvinfer1::DataType::kFLOAT; + } + throw std::runtime_error("AudioDecodeRunner: unsupported tensor dtype: " + dtype); +} + +std::string resolveIOTensorName(nvinfer1::ICudaEngine const* engine, std::string const& preferredName, + nvinfer1::TensorIOMode mode, int32_t modeIndex = 0) +{ + for (int32_t i = 0; i < engine->getNbIOTensors(); ++i) + { + auto const* name = engine->getIOTensorName(i); + if (name != nullptr && preferredName == name && engine->getTensorIOMode(name) == mode) + { + return preferredName; + } + } + + int32_t seen = 0; + for (int32_t i = 0; i < engine->getNbIOTensors(); ++i) + { + auto const* name = engine->getIOTensorName(i); + if (name != nullptr && engine->getTensorIOMode(name) == mode) + { + if (seen == modeIndex) + { + LOG_WARNING( + "AudioDecodeRunner: tensor name '%s' is not present in the TensorRT engine; using binding '%s'", + preferredName.c_str(), name); + return name; + } + ++seen; + } + } + + throw std::runtime_error("AudioDecodeRunner: failed to resolve TensorRT I/O tensor name: " + preferredName); +} + +std::string resolveEnginePath(std::string const& engineDir, nlohmann::json const& configJson) +{ + std::string const engineFile = configJson.value("engine_file", std::string{"audio_decode.engine"}); + std::filesystem::path const preferred = std::filesystem::path(engineDir) / engineFile; + if (std::filesystem::exists(preferred)) + { + return preferred.string(); + } + + for (char const* candidate : {"audio_decode.engine", "audio_decode.engine"}) + { + std::filesystem::path const path = std::filesystem::path(engineDir) / candidate; + if (std::filesystem::exists(path)) + { + return path.string(); + } + } + + return preferred.string(); +} + +int64_t tensorBytes(rt::Tensor const& tensor) +{ + return static_cast(tensor.getShape().volume()) * rt::utils::getTypeSize(tensor.getDataType()); +} + +bool copyTensorToDevice(rt::Tensor& dst, rt::Tensor const& src, cudaStream_t stream) +{ + auto const dstBytes = tensorBytes(dst); + auto const srcBytes = tensorBytes(src); + if (dstBytes == srcBytes) + { + CUDA_CHECK(cudaMemcpyAsync( + dst.rawPointer(), src.rawPointer(), static_cast(dstBytes), cudaMemcpyDeviceToDevice, stream)); + return true; + } + + auto const elements = dst.getShape().volume(); + if (src.getShape().volume() != elements) + { + LOG_ERROR("AudioDecodeRunner: copy shape mismatch, dst=%ld src=%ld", elements, src.getShape().volume()); + return false; + } + + auto const srcType = src.getDataType(); + auto const dstType = dst.getDataType(); + if (srcType == nvinfer1::DataType::kHALF && dstType == nvinfer1::DataType::kFLOAT) + { + std::vector hostSrc(static_cast(elements)); + CUDA_CHECK(cudaMemcpyAsync( + hostSrc.data(), src.rawPointer(), static_cast(srcBytes), cudaMemcpyDeviceToHost, stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + std::vector hostDst(static_cast(elements)); + for (int64_t i = 0; i < elements; ++i) + { + hostDst[static_cast(i)] = __half2float(hostSrc[static_cast(i)]); + } + CUDA_CHECK(cudaMemcpyAsync( + dst.rawPointer(), hostDst.data(), static_cast(dstBytes), cudaMemcpyHostToDevice, stream)); + return true; + } + + if (srcType == nvinfer1::DataType::kFLOAT && dstType == nvinfer1::DataType::kHALF) + { + std::vector hostSrc(static_cast(elements)); + CUDA_CHECK(cudaMemcpyAsync( + hostSrc.data(), src.rawPointer(), static_cast(srcBytes), cudaMemcpyDeviceToHost, stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + std::vector hostDst(static_cast(elements)); + for (int64_t i = 0; i < elements; ++i) + { + hostDst[static_cast(i)] = __float2half(hostSrc[static_cast(i)]); + } + CUDA_CHECK(cudaMemcpyAsync( + dst.rawPointer(), hostDst.data(), static_cast(dstBytes), cudaMemcpyHostToDevice, stream)); + return true; + } + + LOG_ERROR("AudioDecodeRunner: unsupported dtype conversion src=%d dst=%d", static_cast(srcType), + static_cast(dstType)); + return false; +} + +} // namespace + +AudioDecodeRunner::AudioDecodeRunner(std::string const& engineDir, cudaStream_t stream) + : mEngineDir(engineDir) +{ + if (!loadConfig()) + { + throw std::runtime_error("AudioDecodeRunner: failed to load config from " + engineDir); + } + if (!loadEngine(stream)) + { + throw std::runtime_error("AudioDecodeRunner: failed to load TensorRT engine from " + engineDir); + } + if (!validateAndFillConfig()) + { + throw std::runtime_error("AudioDecodeRunner: failed to validate config for " + engineDir); + } + if (!allocateBuffers()) + { + throw std::runtime_error("AudioDecodeRunner: failed to allocate buffers for " + engineDir); + } + + LOG_INFO("AudioDecodeRunner loaded from %s (%s -> %s, latents=%s, pixels=%s)", engineDir.c_str(), + mInputName.c_str(), mOutputName.c_str(), mInputShape.formatString().c_str(), + mOutputShape.formatString().c_str()); +} + +bool AudioDecodeRunner::loadConfig() +{ + auto const configPath = mEngineDir + "/config.json"; + std::ifstream configFile(configPath); + if (!configFile) + { + LOG_ERROR("AudioDecodeRunner: failed to open config file: %s", configPath.c_str()); + return false; + } + + try + { + configFile >> mConfigJson; + } + catch (nlohmann::json::parse_error const& e) + { + LOG_ERROR("AudioDecodeRunner: failed to parse config file: %s", e.what()); + return false; + } + return true; +} + +bool AudioDecodeRunner::loadEngine(cudaStream_t stream) +{ + auto const enginePath = resolveEnginePath(mEngineDir, mConfigJson); + if (!std::filesystem::exists(enginePath)) + { + LOG_ERROR("AudioDecodeRunner: engine not found at %s", enginePath.c_str()); + return false; + } + + file_io::MmapReader engineFileReader(enginePath); + mRuntime.reset(nvinfer1::createInferRuntime(gLogger)); + if (!mRuntime) + { + LOG_ERROR("AudioDecodeRunner: failed to create TensorRT runtime for %s", enginePath.c_str()); + return false; + } + + mEngine.reset(mRuntime->deserializeCudaEngine(engineFileReader.getData(), engineFileReader.getSize())); + if (!mEngine) + { + LOG_ERROR("AudioDecodeRunner: failed to deserialize TensorRT engine: %s", enginePath.c_str()); + return false; + } + + mContext.reset(mEngine->createExecutionContext(nvinfer1::ExecutionContextAllocationStrategy::kUSER_MANAGED)); + if (!mContext) + { + LOG_ERROR("AudioDecodeRunner: failed to create TensorRT execution context for %s", enginePath.c_str()); + return false; + } + + if (!mContext->setOptimizationProfileAsync(0, stream)) + { + LOG_ERROR("AudioDecodeRunner: failed to set optimization profile for %s", enginePath.c_str()); + return false; + } + CUDA_CHECK(cudaStreamSynchronize(stream)); + return true; +} + +bool AudioDecodeRunner::validateAndFillConfig() +{ + auto const modelType = mConfigJson.value("model_type", std::string{"cosmos_avae_decode"}); + auto const component = mConfigJson.value("component", std::string{"audio_decode"}); + if (modelType != "cosmos_avae_decode" && component != "audio_decode") + { + LOG_ERROR("AudioDecodeRunner: unexpected model_type=%s component=%s", modelType.c_str(), component.c_str()); + return false; + } + + if (mConfigJson.contains("input_names") && !mConfigJson.at("input_names").empty()) + { + mInputName = mConfigJson.at("input_names").at(0).get(); + } + if (mConfigJson.contains("output_names") && !mConfigJson.at("output_names").empty()) + { + mOutputName = mConfigJson.at("output_names").at(0).get(); + } + + if (!mConfigJson.contains("inputs") || !mConfigJson.at("inputs").contains(mInputName)) + { + LOG_ERROR("AudioDecodeRunner: config is missing input metadata for %s", mInputName.c_str()); + return false; + } + if (!mConfigJson.contains("outputs") || mConfigJson.at("outputs").empty()) + { + LOG_ERROR("AudioDecodeRunner: config is missing output metadata"); + return false; + } + + auto const& inputMeta = mConfigJson.at("inputs").at(mInputName); + auto const& outputMeta = mConfigJson.at("outputs").at(0); + + mInputShape = coordsFromJson(inputMeta.at("shape")); + mOutputShape = coordsFromJson(outputMeta.at("shape")); + mInputType = dataTypeFromTorchString(inputMeta.at("dtype").get()); + mOutputType = dataTypeFromTorchString(outputMeta.at("dtype").get()); + + if (mInputShape.getNumDims() != 3) + { + LOG_ERROR("AudioDecodeRunner: expected rank-3 sound_latents input (got %d dims)", mInputShape.getNumDims()); + return false; + } + if (mOutputShape.getNumDims() != 2 && mOutputShape.getNumDims() != 3) + { + LOG_ERROR("AudioDecodeRunner: expected rank-2/3 waveform output (got %d dims)", mOutputShape.getNumDims()); + return false; + } + + mInputName = resolveIOTensorName(mEngine.get(), mInputName, nvinfer1::TensorIOMode::kINPUT, 0); + mOutputName = resolveIOTensorName(mEngine.get(), mOutputName, nvinfer1::TensorIOMode::kOUTPUT, 0); + return true; +} + +int64_t AudioDecodeRunner::getRequiredContextMemorySize() const +{ + return mEngine ? mEngine->getDeviceMemorySizeV2() : 0; +} + +bool AudioDecodeRunner::allocateBuffers() +{ + mSoundLatentsTensor = rt::Tensor(mInputShape, rt::DeviceType::kGPU, mInputType, mInputName); + mWaveformTensor = rt::Tensor(mOutputShape, rt::DeviceType::kGPU, mOutputType, mOutputName); + return bindTensors(); +} + +bool AudioDecodeRunner::bindTensors() noexcept +{ + if (!mContext->setInputShape(mInputName.c_str(), mSoundLatentsTensor.getShape().getTRTDims())) + { + LOG_ERROR("AudioDecodeRunner: failed to set input shape for %s", mInputName.c_str()); + return false; + } + if (!mContext->setTensorAddress(mInputName.c_str(), mSoundLatentsTensor.rawPointer())) + { + LOG_ERROR("AudioDecodeRunner: failed to bind input tensor %s", mInputName.c_str()); + return false; + } + if (!mContext->setTensorAddress(mOutputName.c_str(), mWaveformTensor.rawPointer())) + { + LOG_ERROR("AudioDecodeRunner: failed to bind output tensor %s", mOutputName.c_str()); + return false; + } + return true; +} + +bool AudioDecodeRunner::setContextMemory(rt::Tensor& sharedContextMemory) +{ + if (!mEngine) + { + return true; + } + + int64_t const requiredSize = getRequiredContextMemorySize(); + if (sharedContextMemory.getMemoryCapacity() < requiredSize) + { + LOG_ERROR("AudioDecodeRunner: shared context memory (%zu bytes) is smaller than required (%zu bytes)", + static_cast(sharedContextMemory.getMemoryCapacity()), static_cast(requiredSize)); + return false; + } + + mContext->setDeviceMemoryV2(sharedContextMemory.rawPointer(), sharedContextMemory.getMemoryCapacity()); + return true; +} + +bool AudioDecodeRunner::copySoundLatentsFrom(rt::Tensor const& src, cudaStream_t stream) +{ + if (src.getShape() != mSoundLatentsTensor.getShape()) + { + LOG_ERROR("AudioDecodeRunner: latent shape %s does not match engine input %s", + src.getShape().formatString().c_str(), mSoundLatentsTensor.getShape().formatString().c_str()); + return false; + } + return copyTensorToDevice(mSoundLatentsTensor, src, stream); +} + +bool AudioDecodeRunner::copyWaveformTo(rt::Tensor& dst, cudaStream_t stream) const +{ + if (dst.getShape() != mWaveformTensor.getShape()) + { + LOG_ERROR("AudioDecodeRunner: waveform shape %s does not match engine output %s", + dst.getShape().formatString().c_str(), mWaveformTensor.getShape().formatString().c_str()); + return false; + } + return copyTensorToDevice(dst, mWaveformTensor, stream); +} + +bool AudioDecodeRunner::copyWaveformTo(AudioBuffer& audio, cudaStream_t stream) const +{ + if (!audio.buffer) + { + LOG_ERROR("AudioDecodeRunner: AudioBuffer is missing a tensor."); + return false; + } + return copyWaveformTo(*audio.buffer, stream); +} + +bool AudioDecodeRunner::decode(cudaStream_t stream) noexcept +{ + if (!bindTensors()) + { + return false; + } + if (!mContext->enqueueV3(stream)) + { + LOG_ERROR("AudioDecodeRunner: enqueueV3 failed."); + return false; + } + return true; +} + +} // namespace rt +} // namespace trt_edgellm diff --git a/cpp/runtime/audioDecodeRunner.h b/cpp/runtime/audioDecodeRunner.h new file mode 100644 index 00000000..f8ab0e59 --- /dev/null +++ b/cpp/runtime/audioDecodeRunner.h @@ -0,0 +1,124 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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. + */ + +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "common/tensor.h" +#include "runtime/wfmRuntimeUtils.h" + +#include +#include +#include + +#include +#include +#include + +namespace trt_edgellm +{ +namespace rt +{ + +//! AVAE decode engine: sound_latents [B,C,T'] -> waveform [B,1,T] or [B,T]. +class AudioDecodeRunner +{ +public: + //! \p engineDir is the ``audio_decode/`` component directory (contains config.json + engine). + explicit AudioDecodeRunner(std::string const& engineDir, cudaStream_t stream); + + ~AudioDecodeRunner() noexcept = default; + + int64_t getRequiredContextMemorySize() const; + bool setContextMemory(rt::Tensor& sharedContextMemory); + + std::string const& getInputName() const noexcept + { + return mInputName; + } + + std::string const& getOutputName() const noexcept + { + return mOutputName; + } + + rt::Coords const& getSoundLatentShape() const noexcept + { + return mInputShape; + } + + rt::Coords const& getWaveformShape() const noexcept + { + return mOutputShape; + } + + rt::Tensor& getSoundLatents() noexcept + { + return mSoundLatentsTensor; + } + + rt::Tensor const& getSoundLatents() const noexcept + { + return mSoundLatentsTensor; + } + + rt::Tensor& getWaveform() noexcept + { + return mWaveformTensor; + } + + rt::Tensor const& getWaveform() const noexcept + { + return mWaveformTensor; + } + + bool copySoundLatentsFrom(rt::Tensor const& src, cudaStream_t stream); + bool copyWaveformTo(rt::Tensor& dst, cudaStream_t stream) const; + bool copyWaveformTo(AudioBuffer& audio, cudaStream_t stream) const; + bool decode(cudaStream_t stream) noexcept; + +private: + bool loadConfig(); + bool loadEngine(cudaStream_t stream); + bool validateAndFillConfig(); + bool allocateBuffers(); + bool bindTensors() noexcept; + + std::string mEngineDir; + nlohmann::json mConfigJson; + + std::unique_ptr mRuntime; + std::unique_ptr mEngine; + std::unique_ptr mContext; + + std::string mInputName{"sound_latents"}; + std::string mOutputName{"waveform"}; + rt::Coords mInputShape; + rt::Coords mOutputShape; + nvinfer1::DataType mInputType{nvinfer1::DataType::kHALF}; + nvinfer1::DataType mOutputType{nvinfer1::DataType::kHALF}; + + rt::Tensor mSoundLatentsTensor; + rt::Tensor mWaveformTensor; +}; + +} // namespace rt +} // namespace trt_edgellm diff --git a/cpp/runtime/audioEncodeRunner.cpp b/cpp/runtime/audioEncodeRunner.cpp new file mode 100644 index 00000000..6b2e46de --- /dev/null +++ b/cpp/runtime/audioEncodeRunner.cpp @@ -0,0 +1,411 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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. + */ + +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/audioEncodeRunner.h" + +#include "common/checkMacros.h" +#include "common/cudaUtils.h" +#include "common/logger.h" +#include "common/mmapReader.h" + +#include +#include +#include +#include + +namespace trt_edgellm +{ +namespace rt +{ +namespace +{ + +rt::Coords coordsFromJson(nlohmann::json const& shape) +{ + std::vector dims; + dims.reserve(shape.size()); + for (auto const& dim : shape) + { + dims.push_back(dim.get()); + } + return rt::Coords(dims); +} + +nvinfer1::DataType dataTypeFromTorchString(std::string const& dtype) +{ + if (dtype == "torch.float16" || dtype == "float16" || dtype == "fp16") + { + return nvinfer1::DataType::kHALF; + } + if (dtype == "torch.float32" || dtype == "float32" || dtype == "fp32") + { + return nvinfer1::DataType::kFLOAT; + } + throw std::runtime_error("AudioEncodeRunner: unsupported tensor dtype: " + dtype); +} + +std::string resolveIOTensorName(nvinfer1::ICudaEngine const* engine, std::string const& preferredName, + nvinfer1::TensorIOMode mode, int32_t modeIndex = 0) +{ + for (int32_t i = 0; i < engine->getNbIOTensors(); ++i) + { + auto const* name = engine->getIOTensorName(i); + if (name != nullptr && preferredName == name && engine->getTensorIOMode(name) == mode) + { + return preferredName; + } + } + + int32_t seen = 0; + for (int32_t i = 0; i < engine->getNbIOTensors(); ++i) + { + auto const* name = engine->getIOTensorName(i); + if (name != nullptr && engine->getTensorIOMode(name) == mode) + { + if (seen == modeIndex) + { + LOG_WARNING( + "AudioEncodeRunner: tensor name '%s' is not present in the TensorRT engine; using binding '%s'", + preferredName.c_str(), name); + return name; + } + ++seen; + } + } + + throw std::runtime_error("AudioEncodeRunner: failed to resolve TensorRT I/O tensor name: " + preferredName); +} + +std::string resolveEnginePath(std::string const& engineDir, nlohmann::json const& configJson) +{ + std::string const engineFile = configJson.value("engine_file", std::string{"audio_encode.engine"}); + std::filesystem::path const preferred = std::filesystem::path(engineDir) / engineFile; + if (std::filesystem::exists(preferred)) + { + return preferred.string(); + } + + for (char const* candidate : {"audio_encode.engine", "audio_encode.engine"}) + { + std::filesystem::path const path = std::filesystem::path(engineDir) / candidate; + if (std::filesystem::exists(path)) + { + return path.string(); + } + } + + return preferred.string(); +} + +int64_t tensorBytes(rt::Tensor const& tensor) +{ + return static_cast(tensor.getShape().volume()) * rt::utils::getTypeSize(tensor.getDataType()); +} + +bool copyTensorToDevice(rt::Tensor& dst, rt::Tensor const& src, cudaStream_t stream) +{ + auto const dstBytes = tensorBytes(dst); + auto const srcBytes = tensorBytes(src); + if (dstBytes == srcBytes) + { + CUDA_CHECK(cudaMemcpyAsync( + dst.rawPointer(), src.rawPointer(), static_cast(dstBytes), cudaMemcpyDeviceToDevice, stream)); + return true; + } + + auto const elements = dst.getShape().volume(); + if (src.getShape().volume() != elements) + { + LOG_ERROR("AudioEncodeRunner: copy shape mismatch, dst=%ld src=%ld", elements, src.getShape().volume()); + return false; + } + + auto const srcType = src.getDataType(); + auto const dstType = dst.getDataType(); + if (srcType == nvinfer1::DataType::kHALF && dstType == nvinfer1::DataType::kFLOAT) + { + std::vector hostSrc(static_cast(elements)); + CUDA_CHECK(cudaMemcpyAsync( + hostSrc.data(), src.rawPointer(), static_cast(srcBytes), cudaMemcpyDeviceToHost, stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + std::vector hostDst(static_cast(elements)); + for (int64_t i = 0; i < elements; ++i) + { + hostDst[static_cast(i)] = __half2float(hostSrc[static_cast(i)]); + } + CUDA_CHECK(cudaMemcpyAsync( + dst.rawPointer(), hostDst.data(), static_cast(dstBytes), cudaMemcpyHostToDevice, stream)); + return true; + } + + if (srcType == nvinfer1::DataType::kFLOAT && dstType == nvinfer1::DataType::kHALF) + { + std::vector hostSrc(static_cast(elements)); + CUDA_CHECK(cudaMemcpyAsync( + hostSrc.data(), src.rawPointer(), static_cast(srcBytes), cudaMemcpyDeviceToHost, stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + std::vector hostDst(static_cast(elements)); + for (int64_t i = 0; i < elements; ++i) + { + hostDst[static_cast(i)] = __float2half(hostSrc[static_cast(i)]); + } + CUDA_CHECK(cudaMemcpyAsync( + dst.rawPointer(), hostDst.data(), static_cast(dstBytes), cudaMemcpyHostToDevice, stream)); + return true; + } + + LOG_ERROR("AudioEncodeRunner: unsupported dtype conversion src=%d dst=%d", static_cast(srcType), + static_cast(dstType)); + return false; +} + +} // namespace + +AudioEncodeRunner::AudioEncodeRunner(std::string const& engineDir, cudaStream_t stream) + : mEngineDir(engineDir) +{ + if (!loadConfig()) + { + throw std::runtime_error("AudioEncodeRunner: failed to load config from " + engineDir); + } + if (!loadEngine(stream)) + { + throw std::runtime_error("AudioEncodeRunner: failed to load TensorRT engine from " + engineDir); + } + if (!validateAndFillConfig()) + { + throw std::runtime_error("AudioEncodeRunner: failed to validate config for " + engineDir); + } + if (!allocateBuffers()) + { + throw std::runtime_error("AudioEncodeRunner: failed to allocate buffers for " + engineDir); + } + + LOG_INFO("AudioEncodeRunner loaded from %s (%s -> %s, waveform=%s, sound_latents=%s)", engineDir.c_str(), + mInputName.c_str(), mOutputName.c_str(), mInputShape.formatString().c_str(), + mOutputShape.formatString().c_str()); +} + +bool AudioEncodeRunner::loadConfig() +{ + auto const configPath = mEngineDir + "/config.json"; + std::ifstream configFile(configPath); + if (!configFile) + { + LOG_ERROR("AudioEncodeRunner: failed to open config file: %s", configPath.c_str()); + return false; + } + + try + { + configFile >> mConfigJson; + } + catch (nlohmann::json::parse_error const& e) + { + LOG_ERROR("AudioEncodeRunner: failed to parse config file: %s", e.what()); + return false; + } + return true; +} + +bool AudioEncodeRunner::loadEngine(cudaStream_t stream) +{ + auto const enginePath = resolveEnginePath(mEngineDir, mConfigJson); + if (!std::filesystem::exists(enginePath)) + { + LOG_ERROR("AudioEncodeRunner: engine not found at %s", enginePath.c_str()); + return false; + } + + file_io::MmapReader engineFileReader(enginePath); + mRuntime.reset(nvinfer1::createInferRuntime(gLogger)); + if (!mRuntime) + { + LOG_ERROR("AudioEncodeRunner: failed to create TensorRT runtime for %s", enginePath.c_str()); + return false; + } + + mEngine.reset(mRuntime->deserializeCudaEngine(engineFileReader.getData(), engineFileReader.getSize())); + if (!mEngine) + { + LOG_ERROR("AudioEncodeRunner: failed to deserialize TensorRT engine: %s", enginePath.c_str()); + return false; + } + + mContext.reset(mEngine->createExecutionContext(nvinfer1::ExecutionContextAllocationStrategy::kUSER_MANAGED)); + if (!mContext) + { + LOG_ERROR("AudioEncodeRunner: failed to create TensorRT execution context for %s", enginePath.c_str()); + return false; + } + + if (!mContext->setOptimizationProfileAsync(0, stream)) + { + LOG_ERROR("AudioEncodeRunner: failed to set optimization profile for %s", enginePath.c_str()); + return false; + } + CUDA_CHECK(cudaStreamSynchronize(stream)); + return true; +} + +bool AudioEncodeRunner::validateAndFillConfig() +{ + auto const modelType = mConfigJson.value("model_type", std::string{"cosmos_avae_encode"}); + auto const component = mConfigJson.value("component", std::string{"audio_encode"}); + if (modelType != "cosmos_avae_encode" && component != "audio_encode") + { + LOG_ERROR("AudioEncodeRunner: unexpected model_type=%s component=%s", modelType.c_str(), component.c_str()); + return false; + } + + if (mConfigJson.contains("input_names") && !mConfigJson.at("input_names").empty()) + { + mInputName = mConfigJson.at("input_names").at(0).get(); + } + if (mConfigJson.contains("output_names") && !mConfigJson.at("output_names").empty()) + { + mOutputName = mConfigJson.at("output_names").at(0).get(); + } + + if (!mConfigJson.contains("inputs") || !mConfigJson.at("inputs").contains(mInputName)) + { + LOG_ERROR("AudioEncodeRunner: config is missing input metadata for %s", mInputName.c_str()); + return false; + } + if (!mConfigJson.contains("outputs") || mConfigJson.at("outputs").empty()) + { + LOG_ERROR("AudioEncodeRunner: config is missing output metadata"); + return false; + } + + auto const& inputMeta = mConfigJson.at("inputs").at(mInputName); + auto const& outputMeta = mConfigJson.at("outputs").at(0); + + mInputShape = coordsFromJson(inputMeta.at("shape")); + mOutputShape = coordsFromJson(outputMeta.at("shape")); + mInputType = dataTypeFromTorchString(inputMeta.at("dtype").get()); + mOutputType = dataTypeFromTorchString(outputMeta.at("dtype").get()); + + if (mInputShape.getNumDims() != 2 && mInputShape.getNumDims() != 3) + { + LOG_ERROR("AudioEncodeRunner: expected rank-2/3 waveform input (got %d dims)", mInputShape.getNumDims()); + return false; + } + if (mOutputShape.getNumDims() != 3) + { + LOG_ERROR("AudioEncodeRunner: expected rank-3 sound_latents output (got %d dims)", mOutputShape.getNumDims()); + return false; + } + + mInputName = resolveIOTensorName(mEngine.get(), mInputName, nvinfer1::TensorIOMode::kINPUT, 0); + mOutputName = resolveIOTensorName(mEngine.get(), mOutputName, nvinfer1::TensorIOMode::kOUTPUT, 0); + return true; +} + +int64_t AudioEncodeRunner::getRequiredContextMemorySize() const +{ + return mEngine ? mEngine->getDeviceMemorySizeV2() : 0; +} + +bool AudioEncodeRunner::allocateBuffers() +{ + mWaveformTensor = rt::Tensor(mInputShape, rt::DeviceType::kGPU, mInputType, mInputName); + mSoundLatentsTensor = rt::Tensor(mOutputShape, rt::DeviceType::kGPU, mOutputType, mOutputName); + return bindTensors(); +} + +bool AudioEncodeRunner::bindTensors() noexcept +{ + if (!mContext->setInputShape(mInputName.c_str(), mWaveformTensor.getShape().getTRTDims())) + { + LOG_ERROR("AudioEncodeRunner: failed to set input shape for %s", mInputName.c_str()); + return false; + } + if (!mContext->setTensorAddress(mInputName.c_str(), mWaveformTensor.rawPointer())) + { + LOG_ERROR("AudioEncodeRunner: failed to bind input tensor %s", mInputName.c_str()); + return false; + } + if (!mContext->setTensorAddress(mOutputName.c_str(), mSoundLatentsTensor.rawPointer())) + { + LOG_ERROR("AudioEncodeRunner: failed to bind output tensor %s", mOutputName.c_str()); + return false; + } + return true; +} + +bool AudioEncodeRunner::setContextMemory(rt::Tensor& sharedContextMemory) +{ + if (!mEngine) + { + return true; + } + + int64_t const requiredSize = getRequiredContextMemorySize(); + if (sharedContextMemory.getMemoryCapacity() < requiredSize) + { + LOG_ERROR("AudioEncodeRunner: shared context memory (%zu bytes) is smaller than required (%zu bytes)", + static_cast(sharedContextMemory.getMemoryCapacity()), static_cast(requiredSize)); + return false; + } + + mContext->setDeviceMemoryV2(sharedContextMemory.rawPointer(), sharedContextMemory.getMemoryCapacity()); + return true; +} + +bool AudioEncodeRunner::copyWaveformFrom(rt::Tensor const& src, cudaStream_t stream) +{ + if (src.getShape() != mWaveformTensor.getShape()) + { + LOG_ERROR("AudioEncodeRunner: waveform shape %s does not match engine input %s", + src.getShape().formatString().c_str(), mWaveformTensor.getShape().formatString().c_str()); + return false; + } + return copyTensorToDevice(mWaveformTensor, src, stream); +} + +bool AudioEncodeRunner::copyWaveformFrom(AudioBuffer const& audio, cudaStream_t stream) +{ + if (!audio.buffer) + { + LOG_ERROR("AudioEncodeRunner: AudioBuffer is missing a tensor."); + return false; + } + return copyWaveformFrom(*audio.buffer, stream); +} + +bool AudioEncodeRunner::encode(cudaStream_t stream) noexcept +{ + if (!bindTensors()) + { + return false; + } + if (!mContext->enqueueV3(stream)) + { + LOG_ERROR("AudioEncodeRunner: enqueueV3 failed."); + return false; + } + return true; +} + +} // namespace rt +} // namespace trt_edgellm diff --git a/cpp/runtime/audioEncodeRunner.h b/cpp/runtime/audioEncodeRunner.h new file mode 100644 index 00000000..b7ada3c4 --- /dev/null +++ b/cpp/runtime/audioEncodeRunner.h @@ -0,0 +1,123 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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. + */ + +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "common/tensor.h" +#include "runtime/wfmRuntimeUtils.h" + +#include +#include +#include + +#include +#include +#include + +namespace trt_edgellm +{ +namespace rt +{ + +//! AVAE encode engine: waveform [B,1,T] or [B,T] -> sound_latents [B,C,T']. +class AudioEncodeRunner +{ +public: + //! \p engineDir is the ``audio_encode/`` component directory (contains config.json + engine). + explicit AudioEncodeRunner(std::string const& engineDir, cudaStream_t stream); + + ~AudioEncodeRunner() noexcept = default; + + int64_t getRequiredContextMemorySize() const; + bool setContextMemory(rt::Tensor& sharedContextMemory); + + std::string const& getInputName() const noexcept + { + return mInputName; + } + + std::string const& getOutputName() const noexcept + { + return mOutputName; + } + + rt::Coords const& getWaveformShape() const noexcept + { + return mInputShape; + } + + rt::Coords const& getSoundLatentShape() const noexcept + { + return mOutputShape; + } + + rt::Tensor& getWaveform() noexcept + { + return mWaveformTensor; + } + + rt::Tensor const& getWaveform() const noexcept + { + return mWaveformTensor; + } + + rt::Tensor& getSoundLatents() noexcept + { + return mSoundLatentsTensor; + } + + rt::Tensor const& getSoundLatents() const noexcept + { + return mSoundLatentsTensor; + } + + bool copyWaveformFrom(rt::Tensor const& src, cudaStream_t stream); + bool copyWaveformFrom(AudioBuffer const& audio, cudaStream_t stream); + bool encode(cudaStream_t stream) noexcept; + +private: + bool loadConfig(); + bool loadEngine(cudaStream_t stream); + bool validateAndFillConfig(); + bool allocateBuffers(); + bool bindTensors() noexcept; + + std::string mEngineDir; + nlohmann::json mConfigJson; + + std::unique_ptr mRuntime; + std::unique_ptr mEngine; + std::unique_ptr mContext; + + std::string mInputName{"waveform"}; + std::string mOutputName{"sound_latents"}; + rt::Coords mInputShape; + rt::Coords mOutputShape; + nvinfer1::DataType mInputType{nvinfer1::DataType::kHALF}; + nvinfer1::DataType mOutputType{nvinfer1::DataType::kHALF}; + + rt::Tensor mWaveformTensor; + rt::Tensor mSoundLatentsTensor; +}; + +} // namespace rt +} // namespace trt_edgellm diff --git a/cpp/runtime/cosmosDenoiseRunner.cpp b/cpp/runtime/cosmosDenoiseRunner.cpp new file mode 100644 index 00000000..70ca57da --- /dev/null +++ b/cpp/runtime/cosmosDenoiseRunner.cpp @@ -0,0 +1,268 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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. + */ + +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/cosmosDenoiseRunner.h" + +#include "common/logger.h" + +#include +#include +#include + +namespace trt_edgellm +{ +namespace rt +{ +namespace +{ + +std::string componentDir(std::filesystem::path const& engineRoot, char const* name) +{ + return (engineRoot / name).string(); +} + +bool componentExists(std::filesystem::path const& engineRoot, char const* name) +{ + return std::filesystem::exists(engineRoot / name / "config.json"); +} + +} // namespace + +CosmosDenoiseRunner::CosmosDenoiseRunner( + std::filesystem::path const& engineRoot, CosmosEngineConfig const& config, cudaStream_t stream) + : mConfig(config) +{ + if (!componentExists(engineRoot, "embed") || !componentExists(engineRoot, "mot_backbone") + || !componentExists(engineRoot, "denoise_head")) + { + throw std::runtime_error( + "CosmosDenoiseRunner: missing embed/, mot_backbone/, or denoise_head/ under " + engineRoot.string()); + } + + mEmbedRunner = std::make_unique(componentDir(engineRoot, "embed"), stream); + mMotBackboneRunner = std::make_unique(componentDir(engineRoot, "mot_backbone"), stream); + mVisionHeadRunner = std::make_unique(componentDir(engineRoot, "denoise_head"), stream); + + if (componentExists(engineRoot, "denoise_head_sound")) + { + mSoundHeadRunner = std::make_unique(componentDir(engineRoot, "denoise_head_sound"), stream); + LOG_INFO("CosmosDenoiseRunner: loaded optional denoise_head_sound."); + } + + LOG_INFO("CosmosDenoiseRunner loaded from %s (embed + mot_backbone + denoise_head%s).", engineRoot.string().c_str(), + mSoundHeadRunner ? " + denoise_head_sound" : ""); +} + +bool CosmosDenoiseRunner::isReady() const noexcept +{ + return mEmbedRunner && mMotBackboneRunner && mVisionHeadRunner; +} + +bool CosmosDenoiseRunner::hasSoundPath() const noexcept +{ + return static_cast(mSoundHeadRunner); +} + +int64_t CosmosDenoiseRunner::getRequiredContextMemorySize() const +{ + if (!isReady()) + { + return 0; + } + + int64_t size = std::max({mEmbedRunner->getRequiredContextMemorySize(), + mMotBackboneRunner->getRequiredContextMemorySize(), mVisionHeadRunner->getRequiredContextMemorySize()}); + if (mSoundHeadRunner) + { + size = std::max(size, mSoundHeadRunner->getRequiredContextMemorySize()); + } + return size; +} + +bool CosmosDenoiseRunner::setContextMemory(rt::Tensor& sharedContextMemory) +{ + if (!isReady()) + { + return true; + } + + bool ok = mEmbedRunner->setContextMemory(sharedContextMemory) + && mMotBackboneRunner->setContextMemory(sharedContextMemory) + && mVisionHeadRunner->setContextMemory(sharedContextMemory); + if (mSoundHeadRunner) + { + ok = ok && mSoundHeadRunner->setContextMemory(sharedContextMemory); + } + return ok; +} + +bool CosmosDenoiseRunner::runDenoiseStep(CosmosTextPhase0 const& phase0, CosmosDenoiseLatents latents, + CosmosFlowScheduler& visionScheduler, CosmosFlowScheduler* soundScheduler, cudaStream_t stream) +{ + if (latents.vision == nullptr) + { + LOG_ERROR("CosmosDenoiseRunner: vision latents are required."); + return false; + } + + float const timestep = visionScheduler.getTimestepValue(); + + if (!mEmbedRunner->copyLatentsFrom(*latents.vision, stream)) + { + LOG_ERROR("CosmosDenoiseRunner: embed copyLatentsFrom failed."); + return false; + } + if (!mEmbedRunner->setTimestep(timestep, stream)) + { + LOG_ERROR("CosmosDenoiseRunner: embed setTimestep failed."); + return false; + } + if (!mEmbedRunner->embed(stream)) + { + LOG_ERROR("CosmosDenoiseRunner: embed TRT failed."); + return false; + } + + if (!mMotBackboneRunner->copyUndSeqFrom(phase0.undSeq, stream)) + { + LOG_ERROR("CosmosDenoiseRunner: backbone copyUndSeqFrom failed."); + return false; + } + if (!mMotBackboneRunner->copyGenSeqFrom(mEmbedRunner->getGenSeq(), stream)) + { + LOG_ERROR("CosmosDenoiseRunner: backbone copyGenSeqFrom failed."); + return false; + } + if (!mMotBackboneRunner->copyRotaryFrom(phase0, stream)) + { + LOG_ERROR("CosmosDenoiseRunner: backbone copyRotaryFrom failed."); + return false; + } + if (!mMotBackboneRunner->runBackbone(stream)) + { + LOG_ERROR("CosmosDenoiseRunner: mot_backbone TRT failed."); + return false; + } + + auto const& lastHidden = mMotBackboneRunner->getLastHiddenState(); + + if (!mVisionHeadRunner->copyLastHiddenFrom(lastHidden, stream)) + { + LOG_ERROR("CosmosDenoiseRunner: vision head copyLastHiddenFrom failed."); + return false; + } + if (!mVisionHeadRunner->runHead(stream)) + { + LOG_ERROR("CosmosDenoiseRunner: vision head TRT failed."); + return false; + } + if (!visionScheduler.stepLatents(*latents.vision, mVisionHeadRunner->getPredLatents(), stream)) + { + LOG_ERROR("CosmosDenoiseRunner: vision scheduler step failed."); + return false; + } + + if (latents.sound != nullptr && mSoundHeadRunner && soundScheduler != nullptr) + { + if (!mSoundHeadRunner->copyLastHiddenFrom(lastHidden, stream)) + { + LOG_ERROR("CosmosDenoiseRunner: sound head copyLastHiddenFrom failed."); + return false; + } + if (!mSoundHeadRunner->runHead(stream)) + { + LOG_ERROR("CosmosDenoiseRunner: sound head TRT failed."); + return false; + } + if (!soundScheduler->stepLatents(*latents.sound, mSoundHeadRunner->getPredLatents(), stream)) + { + LOG_ERROR("CosmosDenoiseRunner: sound scheduler step failed."); + return false; + } + } + + return true; +} + +bool CosmosDenoiseRunner::sampleLatents( + CosmosTextPhase0 const& phase0, CosmosDenoiseLatents latents, int32_t numInferenceSteps, cudaStream_t stream) +{ + if (!isReady()) + { + LOG_ERROR("CosmosDenoiseRunner: denoise engines are not loaded."); + return false; + } + + if (latents.vision == nullptr) + { + LOG_ERROR("CosmosDenoiseRunner: vision latents pointer is null."); + return false; + } + + if (latents.sound != nullptr && !mSoundHeadRunner) + { + LOG_ERROR("CosmosDenoiseRunner: sound latents provided but denoise_head_sound is not loaded."); + return false; + } + + if (numInferenceSteps <= 0) + { + LOG_ERROR("CosmosDenoiseRunner: numInferenceSteps must be positive."); + return false; + } + + CosmosFlowScheduler visionScheduler(mConfig); + visionScheduler.setTimesteps(numInferenceSteps); + + std::unique_ptr soundScheduler; + if (latents.sound != nullptr && mSoundHeadRunner) + { + soundScheduler = std::make_unique(mConfig); + soundScheduler->setTimesteps(numInferenceSteps); + } + + LOG_INFO("CosmosDenoiseRunner: starting denoise loop (%d steps, sound=%s).", numInferenceSteps, + latents.sound != nullptr ? "yes" : "no"); + + for (int32_t step = 0; step < numInferenceSteps; ++step) + { + if (!runDenoiseStep(phase0, latents, visionScheduler, soundScheduler.get(), stream)) + { + LOG_ERROR("CosmosDenoiseRunner: denoise step %d failed.", step); + return false; + } + } + + LOG_INFO("CosmosDenoiseRunner: denoise loop complete."); + return true; +} + +bool CosmosDenoiseRunner::sampleVisionLatents( + CosmosTextPhase0 const& phase0, rt::Tensor& visionLatents, int32_t numInferenceSteps, cudaStream_t stream) +{ + CosmosDenoiseLatents latents{}; + latents.vision = &visionLatents; + return sampleLatents(phase0, latents, numInferenceSteps, stream); +} + +} // namespace rt +} // namespace trt_edgellm diff --git a/cpp/runtime/cosmosDenoiseRunner.h b/cpp/runtime/cosmosDenoiseRunner.h new file mode 100644 index 00000000..d2423061 --- /dev/null +++ b/cpp/runtime/cosmosDenoiseRunner.h @@ -0,0 +1,79 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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. + */ + +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "runtime/denoiseHeadRunner.h" +#include "runtime/embedRunner.h" +#include "runtime/motBackboneRunner.h" +#include "runtime/wfmRuntimeUtils.h" + +#include +#include +#include + +namespace trt_edgellm +{ +namespace rt +{ + +//! Orchestrates Cosmos denoise TRT engines + scheduler (ActionRunner analogue for WFM). +//! +//! Edge today: vision-only path (embed -> backbone -> head -> UniPC step). +//! Omni: optional sound head + independent sound scheduler; shared backbone per step. +class CosmosDenoiseRunner +{ +public: + //! \p engineRoot contains ``embed/``, ``mot_backbone/``, ``denoise_head/`` subdirs. + //! Optional ``denoise_head_sound/`` enables the sound denoise path. + explicit CosmosDenoiseRunner( + std::filesystem::path const& engineRoot, CosmosEngineConfig const& config, cudaStream_t stream); + + ~CosmosDenoiseRunner() noexcept = default; + + bool isReady() const noexcept; + bool hasSoundPath() const noexcept; + + int64_t getRequiredContextMemorySize() const; + bool setContextMemory(rt::Tensor& sharedContextMemory); + + //! Run the denoise scheduler loop; updates vision and/or sound latents in place. + bool sampleLatents( + CosmosTextPhase0 const& phase0, CosmosDenoiseLatents latents, int32_t numInferenceSteps, cudaStream_t stream); + + //! Vision-only convenience wrapper. + bool sampleVisionLatents( + CosmosTextPhase0 const& phase0, rt::Tensor& visionLatents, int32_t numInferenceSteps, cudaStream_t stream); + +private: + bool runDenoiseStep(CosmosTextPhase0 const& phase0, CosmosDenoiseLatents latents, + CosmosFlowScheduler& visionScheduler, CosmosFlowScheduler* soundScheduler, cudaStream_t stream); + + CosmosEngineConfig mConfig{}; + std::unique_ptr mEmbedRunner; + std::unique_ptr mMotBackboneRunner; + std::unique_ptr mVisionHeadRunner; + std::unique_ptr mSoundHeadRunner; +}; + +} // namespace rt +} // namespace trt_edgellm diff --git a/cpp/runtime/denoiseHeadRunner.cpp b/cpp/runtime/denoiseHeadRunner.cpp new file mode 100644 index 00000000..23aed65a --- /dev/null +++ b/cpp/runtime/denoiseHeadRunner.cpp @@ -0,0 +1,405 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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. + */ + +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/denoiseHeadRunner.h" + +#include "common/checkMacros.h" +#include "common/cudaUtils.h" +#include "common/logger.h" +#include "common/mmapReader.h" + +#include +#include +#include +#include + +namespace trt_edgellm +{ +namespace rt +{ +namespace +{ + +rt::Coords coordsFromJson(nlohmann::json const& shape) +{ + std::vector dims; + dims.reserve(shape.size()); + for (auto const& dim : shape) + { + dims.push_back(dim.get()); + } + return rt::Coords(dims); +} + +nvinfer1::DataType dataTypeFromTorchString(std::string const& dtype) +{ + if (dtype == "torch.float16" || dtype == "float16" || dtype == "fp16") + { + return nvinfer1::DataType::kHALF; + } + if (dtype == "torch.float32" || dtype == "float32" || dtype == "fp32") + { + return nvinfer1::DataType::kFLOAT; + } + throw std::runtime_error("DenoiseHeadRunner: unsupported tensor dtype: " + dtype); +} + +std::string resolveIOTensorName(nvinfer1::ICudaEngine const* engine, std::string const& preferredName, + nvinfer1::TensorIOMode mode, int32_t modeIndex = 0) +{ + for (int32_t i = 0; i < engine->getNbIOTensors(); ++i) + { + auto const* name = engine->getIOTensorName(i); + if (name != nullptr && preferredName == name && engine->getTensorIOMode(name) == mode) + { + return preferredName; + } + } + + int32_t seen = 0; + for (int32_t i = 0; i < engine->getNbIOTensors(); ++i) + { + auto const* name = engine->getIOTensorName(i); + if (name != nullptr && engine->getTensorIOMode(name) == mode) + { + if (seen == modeIndex) + { + LOG_WARNING( + "DenoiseHeadRunner: tensor name '%s' is not present in the TensorRT engine; using binding '%s'", + preferredName.c_str(), name); + return name; + } + ++seen; + } + } + + throw std::runtime_error("DenoiseHeadRunner: failed to resolve TensorRT I/O tensor name: " + preferredName); +} + +std::string resolveEnginePath(std::string const& engineDir, nlohmann::json const& configJson) +{ + std::string const engineFile = configJson.value("engine_file", std::string{"denoise_head.engine"}); + std::filesystem::path const preferred = std::filesystem::path(engineDir) / engineFile; + if (std::filesystem::exists(preferred)) + { + return preferred.string(); + } + + for (char const* candidate : {"denoise_head.engine", "head.engine"}) + { + std::filesystem::path const path = std::filesystem::path(engineDir) / candidate; + if (std::filesystem::exists(path)) + { + return path.string(); + } + } + + return preferred.string(); +} + +int64_t tensorBytes(rt::Tensor const& tensor) +{ + return static_cast(tensor.getShape().volume()) * rt::utils::getTypeSize(tensor.getDataType()); +} + +bool copyTensorToDevice(rt::Tensor& dst, rt::Tensor const& src, cudaStream_t stream) +{ + auto const dstBytes = tensorBytes(dst); + auto const srcBytes = tensorBytes(src); + if (dstBytes == srcBytes) + { + CUDA_CHECK(cudaMemcpyAsync( + dst.rawPointer(), src.rawPointer(), static_cast(dstBytes), cudaMemcpyDeviceToDevice, stream)); + return true; + } + + auto const elements = dst.getShape().volume(); + if (src.getShape().volume() != elements) + { + LOG_ERROR("DenoiseHeadRunner: copy shape mismatch, dst=%ld src=%ld", elements, src.getShape().volume()); + return false; + } + + auto const srcType = src.getDataType(); + auto const dstType = dst.getDataType(); + if (srcType == nvinfer1::DataType::kHALF && dstType == nvinfer1::DataType::kFLOAT) + { + std::vector hostSrc(static_cast(elements)); + CUDA_CHECK(cudaMemcpyAsync( + hostSrc.data(), src.rawPointer(), static_cast(srcBytes), cudaMemcpyDeviceToHost, stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + std::vector hostDst(static_cast(elements)); + for (int64_t i = 0; i < elements; ++i) + { + hostDst[static_cast(i)] = __half2float(hostSrc[static_cast(i)]); + } + CUDA_CHECK(cudaMemcpyAsync( + dst.rawPointer(), hostDst.data(), static_cast(dstBytes), cudaMemcpyHostToDevice, stream)); + return true; + } + + if (srcType == nvinfer1::DataType::kFLOAT && dstType == nvinfer1::DataType::kHALF) + { + std::vector hostSrc(static_cast(elements)); + CUDA_CHECK(cudaMemcpyAsync( + hostSrc.data(), src.rawPointer(), static_cast(srcBytes), cudaMemcpyDeviceToHost, stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + std::vector hostDst(static_cast(elements)); + for (int64_t i = 0; i < elements; ++i) + { + hostDst[static_cast(i)] = __float2half(hostSrc[static_cast(i)]); + } + CUDA_CHECK(cudaMemcpyAsync( + dst.rawPointer(), hostDst.data(), static_cast(dstBytes), cudaMemcpyHostToDevice, stream)); + return true; + } + + LOG_ERROR("DenoiseHeadRunner: unsupported dtype conversion src=%d dst=%d", static_cast(srcType), + static_cast(dstType)); + return false; +} + +} // namespace + +DenoiseHeadRunner::DenoiseHeadRunner(std::string const& engineDir, cudaStream_t stream) + : mEngineDir(engineDir) +{ + if (!loadConfig()) + { + throw std::runtime_error("DenoiseHeadRunner: failed to load config from " + engineDir); + } + if (!loadEngine(stream)) + { + throw std::runtime_error("DenoiseHeadRunner: failed to load TensorRT engine from " + engineDir); + } + if (!validateAndFillConfig()) + { + throw std::runtime_error("DenoiseHeadRunner: failed to validate config for " + engineDir); + } + if (!allocateBuffers()) + { + throw std::runtime_error("DenoiseHeadRunner: failed to allocate buffers for " + engineDir); + } + + LOG_INFO("DenoiseHeadRunner loaded from %s (%s -> %s, last_hidden=%s, pred_latents=%s)", engineDir.c_str(), + mInputName.c_str(), mOutputName.c_str(), mInputShape.formatString().c_str(), + mOutputShape.formatString().c_str()); +} + +bool DenoiseHeadRunner::loadConfig() +{ + auto const configPath = mEngineDir + "/config.json"; + std::ifstream configFile(configPath); + if (!configFile) + { + LOG_ERROR("DenoiseHeadRunner: failed to open config file: %s", configPath.c_str()); + return false; + } + + try + { + configFile >> mConfigJson; + } + catch (nlohmann::json::parse_error const& e) + { + LOG_ERROR("DenoiseHeadRunner: failed to parse config file: %s", e.what()); + return false; + } + return true; +} + +bool DenoiseHeadRunner::loadEngine(cudaStream_t stream) +{ + auto const enginePath = resolveEnginePath(mEngineDir, mConfigJson); + if (!std::filesystem::exists(enginePath)) + { + LOG_ERROR("DenoiseHeadRunner: engine not found at %s", enginePath.c_str()); + return false; + } + + file_io::MmapReader engineFileReader(enginePath); + mRuntime.reset(nvinfer1::createInferRuntime(gLogger)); + if (!mRuntime) + { + LOG_ERROR("DenoiseHeadRunner: failed to create TensorRT runtime for %s", enginePath.c_str()); + return false; + } + + mEngine.reset(mRuntime->deserializeCudaEngine(engineFileReader.getData(), engineFileReader.getSize())); + if (!mEngine) + { + LOG_ERROR("DenoiseHeadRunner: failed to deserialize TensorRT engine: %s", enginePath.c_str()); + return false; + } + + mContext.reset(mEngine->createExecutionContext(nvinfer1::ExecutionContextAllocationStrategy::kUSER_MANAGED)); + if (!mContext) + { + LOG_ERROR("DenoiseHeadRunner: failed to create TensorRT execution context for %s", enginePath.c_str()); + return false; + } + + if (!mContext->setOptimizationProfileAsync(0, stream)) + { + LOG_ERROR("DenoiseHeadRunner: failed to set optimization profile for %s", enginePath.c_str()); + return false; + } + CUDA_CHECK(cudaStreamSynchronize(stream)); + return true; +} + +bool DenoiseHeadRunner::validateAndFillConfig() +{ + auto const modelType = mConfigJson.value("model_type", std::string{"cosmos_denoise_head"}); + auto const component = mConfigJson.value("component", std::string{"denoise_head"}); + bool const validVisionHead = modelType == "cosmos_denoise_head" || component == "denoise_head"; + bool const validSoundHead = modelType == "cosmos_denoise_head_sound" || component == "denoise_head_sound"; + if (!validVisionHead && !validSoundHead) + { + LOG_ERROR("DenoiseHeadRunner: unexpected model_type=%s component=%s", modelType.c_str(), component.c_str()); + return false; + } + + if (mConfigJson.contains("input_names") && !mConfigJson.at("input_names").empty()) + { + mInputName = mConfigJson.at("input_names").at(0).get(); + } + if (mConfigJson.contains("output_names") && !mConfigJson.at("output_names").empty()) + { + mOutputName = mConfigJson.at("output_names").at(0).get(); + } + + if (!mConfigJson.contains("inputs") || !mConfigJson.at("inputs").contains(mInputName)) + { + LOG_ERROR("DenoiseHeadRunner: config is missing input metadata for %s", mInputName.c_str()); + return false; + } + if (!mConfigJson.contains("outputs") || mConfigJson.at("outputs").empty()) + { + LOG_ERROR("DenoiseHeadRunner: config is missing output metadata"); + return false; + } + + auto const& inputMeta = mConfigJson.at("inputs").at(mInputName); + auto const& outputMeta = mConfigJson.at("outputs").at(0); + + mInputShape = coordsFromJson(inputMeta.at("shape")); + mOutputShape = coordsFromJson(outputMeta.at("shape")); + mInputType = dataTypeFromTorchString(inputMeta.at("dtype").get()); + mOutputType = dataTypeFromTorchString(outputMeta.at("dtype").get()); + + if (mInputShape.getNumDims() != 2) + { + LOG_ERROR("DenoiseHeadRunner: expected rank-2 last_hidden_state input, got %d dims", mInputShape.getNumDims()); + return false; + } + if (mOutputShape.getNumDims() != 3 && mOutputShape.getNumDims() != 5) + { + LOG_ERROR( + "DenoiseHeadRunner: expected rank-3 pred_sound_latents or rank-5 pred_vision_latents output, got %d dims", + mOutputShape.getNumDims()); + return false; + } + + mInputName = resolveIOTensorName(mEngine.get(), mInputName, nvinfer1::TensorIOMode::kINPUT, 0); + mOutputName = resolveIOTensorName(mEngine.get(), mOutputName, nvinfer1::TensorIOMode::kOUTPUT, 0); + return true; +} + +int64_t DenoiseHeadRunner::getRequiredContextMemorySize() const +{ + return mEngine ? mEngine->getDeviceMemorySizeV2() : 0; +} + +bool DenoiseHeadRunner::allocateBuffers() +{ + mLastHiddenTensor = rt::Tensor(mInputShape, rt::DeviceType::kGPU, mInputType, mInputName); + mPredLatentsTensor = rt::Tensor(mOutputShape, rt::DeviceType::kGPU, mOutputType, mOutputName); + return bindTensors(); +} + +bool DenoiseHeadRunner::bindTensors() noexcept +{ + if (!mContext->setInputShape(mInputName.c_str(), mLastHiddenTensor.getShape().getTRTDims())) + { + LOG_ERROR("DenoiseHeadRunner: failed to set input shape for %s", mInputName.c_str()); + return false; + } + if (!mContext->setTensorAddress(mInputName.c_str(), mLastHiddenTensor.rawPointer())) + { + LOG_ERROR("DenoiseHeadRunner: failed to bind input tensor %s", mInputName.c_str()); + return false; + } + if (!mContext->setTensorAddress(mOutputName.c_str(), mPredLatentsTensor.rawPointer())) + { + LOG_ERROR("DenoiseHeadRunner: failed to bind output tensor %s", mOutputName.c_str()); + return false; + } + return true; +} + +bool DenoiseHeadRunner::setContextMemory(rt::Tensor& sharedContextMemory) +{ + if (!mEngine) + { + return true; + } + + int64_t const requiredSize = getRequiredContextMemorySize(); + if (sharedContextMemory.getMemoryCapacity() < requiredSize) + { + LOG_ERROR("DenoiseHeadRunner: shared context memory (%zu bytes) is smaller than required (%zu bytes)", + static_cast(sharedContextMemory.getMemoryCapacity()), static_cast(requiredSize)); + return false; + } + + mContext->setDeviceMemoryV2(sharedContextMemory.rawPointer(), sharedContextMemory.getMemoryCapacity()); + return true; +} + +bool DenoiseHeadRunner::copyLastHiddenFrom(rt::Tensor const& src, cudaStream_t stream) +{ + if (src.getShape() != mLastHiddenTensor.getShape()) + { + LOG_ERROR("DenoiseHeadRunner: last_hidden shape %s does not match engine input %s", + src.getShape().formatString().c_str(), mLastHiddenTensor.getShape().formatString().c_str()); + return false; + } + return copyTensorToDevice(mLastHiddenTensor, src, stream); +} + +bool DenoiseHeadRunner::runHead(cudaStream_t stream) noexcept +{ + if (!bindTensors()) + { + return false; + } + if (!mContext->enqueueV3(stream)) + { + LOG_ERROR("DenoiseHeadRunner: enqueueV3 failed."); + return false; + } + return true; +} + +} // namespace rt +} // namespace trt_edgellm diff --git a/cpp/runtime/denoiseHeadRunner.h b/cpp/runtime/denoiseHeadRunner.h new file mode 100644 index 00000000..8ec8bdd8 --- /dev/null +++ b/cpp/runtime/denoiseHeadRunner.h @@ -0,0 +1,122 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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. + */ + +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "common/tensor.h" +#include "runtime/wfmRuntimeUtils.h" + +#include +#include +#include + +#include +#include +#include + +namespace trt_edgellm +{ +namespace rt +{ + +//! Cosmos vision denoise head: last_hidden_state -> pred vision latents. +class DenoiseHeadRunner +{ +public: + //! \p engineDir is the ``denoise_head/`` component directory (contains config.json + engine). + explicit DenoiseHeadRunner(std::string const& engineDir, cudaStream_t stream); + + ~DenoiseHeadRunner() noexcept = default; + + int64_t getRequiredContextMemorySize() const; + bool setContextMemory(rt::Tensor& sharedContextMemory); + + std::string const& getInputName() const noexcept + { + return mInputName; + } + + std::string const& getOutputName() const noexcept + { + return mOutputName; + } + + rt::Coords const& getLastHiddenShape() const noexcept + { + return mInputShape; + } + + rt::Coords const& getPredLatentShape() const noexcept + { + return mOutputShape; + } + + rt::Tensor& getLastHiddenState() noexcept + { + return mLastHiddenTensor; + } + + rt::Tensor const& getLastHiddenState() const noexcept + { + return mLastHiddenTensor; + } + + rt::Tensor& getPredLatents() noexcept + { + return mPredLatentsTensor; + } + + rt::Tensor const& getPredLatents() const noexcept + { + return mPredLatentsTensor; + } + + bool copyLastHiddenFrom(rt::Tensor const& src, cudaStream_t stream); + bool runHead(cudaStream_t stream) noexcept; + +private: + bool loadConfig(); + bool loadEngine(cudaStream_t stream); + bool validateAndFillConfig(); + bool allocateBuffers(); + bool bindTensors() noexcept; + + std::string mEngineDir; + nlohmann::json mConfigJson; + + std::unique_ptr mRuntime; + std::unique_ptr mEngine; + std::unique_ptr mContext; + + std::string mInputName{"last_hidden_state"}; + std::string mOutputName{"pred_vision_latents"}; + rt::Coords mInputShape; + rt::Coords mOutputShape; + nvinfer1::DataType mInputType{nvinfer1::DataType::kHALF}; + nvinfer1::DataType mOutputType{nvinfer1::DataType::kHALF}; + + rt::Tensor mLastHiddenTensor; + rt::Tensor mPredLatentsTensor; +}; + +} // namespace rt +} // namespace trt_edgellm diff --git a/cpp/runtime/embedRunner.cpp b/cpp/runtime/embedRunner.cpp new file mode 100644 index 00000000..82633e1e --- /dev/null +++ b/cpp/runtime/embedRunner.cpp @@ -0,0 +1,468 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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. + */ + +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/embedRunner.h" + +#include "common/checkMacros.h" +#include "common/cudaUtils.h" +#include "common/logger.h" +#include "common/mmapReader.h" + +#include +#include +#include +#include + +namespace trt_edgellm +{ +namespace rt +{ +namespace +{ + +rt::Coords coordsFromJson(nlohmann::json const& shape) +{ + std::vector dims; + dims.reserve(shape.size()); + for (auto const& dim : shape) + { + dims.push_back(dim.get()); + } + return rt::Coords(dims); +} + +nvinfer1::DataType dataTypeFromTorchString(std::string const& dtype) +{ + if (dtype == "torch.float16" || dtype == "float16" || dtype == "fp16") + { + return nvinfer1::DataType::kHALF; + } + if (dtype == "torch.float32" || dtype == "float32" || dtype == "fp32") + { + return nvinfer1::DataType::kFLOAT; + } + throw std::runtime_error("EmbedRunner: unsupported tensor dtype: " + dtype); +} + +std::string resolveIOTensorName(nvinfer1::ICudaEngine const* engine, std::string const& preferredName, + nvinfer1::TensorIOMode mode, int32_t modeIndex = 0) +{ + for (int32_t i = 0; i < engine->getNbIOTensors(); ++i) + { + auto const* name = engine->getIOTensorName(i); + if (name != nullptr && preferredName == name && engine->getTensorIOMode(name) == mode) + { + return preferredName; + } + } + + int32_t seen = 0; + for (int32_t i = 0; i < engine->getNbIOTensors(); ++i) + { + auto const* name = engine->getIOTensorName(i); + if (name != nullptr && engine->getTensorIOMode(name) == mode) + { + if (seen == modeIndex) + { + LOG_WARNING("EmbedRunner: tensor name '%s' is not present in the TensorRT engine; using binding '%s'", + preferredName.c_str(), name); + return name; + } + ++seen; + } + } + + throw std::runtime_error("EmbedRunner: failed to resolve TensorRT I/O tensor name: " + preferredName); +} + +std::string resolveEnginePath(std::string const& engineDir, nlohmann::json const& configJson) +{ + std::string const engineFile = configJson.value("engine_file", std::string{"embed.engine"}); + std::filesystem::path const preferred = std::filesystem::path(engineDir) / engineFile; + if (std::filesystem::exists(preferred)) + { + return preferred.string(); + } + + for (char const* candidate : {"embed.engine", "gen_embed.engine"}) + { + std::filesystem::path const path = std::filesystem::path(engineDir) / candidate; + if (std::filesystem::exists(path)) + { + return path.string(); + } + } + + return preferred.string(); +} + +int64_t tensorBytes(rt::Tensor const& tensor) +{ + return static_cast(tensor.getShape().volume()) * rt::utils::getTypeSize(tensor.getDataType()); +} + +bool copyTensorToDevice(rt::Tensor& dst, rt::Tensor const& src, cudaStream_t stream) +{ + auto const dstBytes = tensorBytes(dst); + auto const srcBytes = tensorBytes(src); + if (dstBytes == srcBytes) + { + CUDA_CHECK(cudaMemcpyAsync( + dst.rawPointer(), src.rawPointer(), static_cast(dstBytes), cudaMemcpyDeviceToDevice, stream)); + return true; + } + + auto const elements = dst.getShape().volume(); + if (src.getShape().volume() != elements) + { + LOG_ERROR("EmbedRunner: copy shape mismatch, dst=%ld src=%ld", elements, src.getShape().volume()); + return false; + } + + auto const srcType = src.getDataType(); + auto const dstType = dst.getDataType(); + if (srcType == nvinfer1::DataType::kHALF && dstType == nvinfer1::DataType::kFLOAT) + { + std::vector hostSrc(static_cast(elements)); + CUDA_CHECK(cudaMemcpyAsync( + hostSrc.data(), src.rawPointer(), static_cast(srcBytes), cudaMemcpyDeviceToHost, stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + std::vector hostDst(static_cast(elements)); + for (int64_t i = 0; i < elements; ++i) + { + hostDst[static_cast(i)] = __half2float(hostSrc[static_cast(i)]); + } + CUDA_CHECK(cudaMemcpyAsync( + dst.rawPointer(), hostDst.data(), static_cast(dstBytes), cudaMemcpyHostToDevice, stream)); + return true; + } + + if (srcType == nvinfer1::DataType::kFLOAT && dstType == nvinfer1::DataType::kHALF) + { + std::vector hostSrc(static_cast(elements)); + CUDA_CHECK(cudaMemcpyAsync( + hostSrc.data(), src.rawPointer(), static_cast(srcBytes), cudaMemcpyDeviceToHost, stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + std::vector hostDst(static_cast(elements)); + for (int64_t i = 0; i < elements; ++i) + { + hostDst[static_cast(i)] = __float2half(hostSrc[static_cast(i)]); + } + CUDA_CHECK(cudaMemcpyAsync( + dst.rawPointer(), hostDst.data(), static_cast(dstBytes), cudaMemcpyHostToDevice, stream)); + return true; + } + + LOG_ERROR("EmbedRunner: unsupported dtype conversion src=%d dst=%d", static_cast(srcType), + static_cast(dstType)); + return false; +} + +} // namespace + +EmbedRunner::EmbedRunner(std::string const& engineDir, cudaStream_t stream) + : mEngineDir(engineDir) +{ + if (!loadConfig()) + { + throw std::runtime_error("EmbedRunner: failed to load config from " + engineDir); + } + if (!loadEngine(stream)) + { + throw std::runtime_error("EmbedRunner: failed to load TensorRT engine from " + engineDir); + } + if (!validateAndFillConfig()) + { + throw std::runtime_error("EmbedRunner: failed to validate config for " + engineDir); + } + if (!allocateBuffers()) + { + throw std::runtime_error("EmbedRunner: failed to allocate buffers for " + engineDir); + } + + LOG_INFO("EmbedRunner loaded from %s (%s + %s -> %s, latents=%s, gen_seq=%s)", engineDir.c_str(), + mLatentsInputName.c_str(), mTimestepInputName.c_str(), mOutputName.c_str(), + mLatentsShape.formatString().c_str(), mOutputShape.formatString().c_str()); +} + +bool EmbedRunner::loadConfig() +{ + auto const configPath = mEngineDir + "/config.json"; + std::ifstream configFile(configPath); + if (!configFile) + { + LOG_ERROR("EmbedRunner: failed to open config file: %s", configPath.c_str()); + return false; + } + + try + { + configFile >> mConfigJson; + } + catch (nlohmann::json::parse_error const& e) + { + LOG_ERROR("EmbedRunner: failed to parse config file: %s", e.what()); + return false; + } + return true; +} + +bool EmbedRunner::loadEngine(cudaStream_t stream) +{ + auto const enginePath = resolveEnginePath(mEngineDir, mConfigJson); + if (!std::filesystem::exists(enginePath)) + { + LOG_ERROR("EmbedRunner: engine not found at %s", enginePath.c_str()); + return false; + } + + file_io::MmapReader engineFileReader(enginePath); + mRuntime.reset(nvinfer1::createInferRuntime(gLogger)); + if (!mRuntime) + { + LOG_ERROR("EmbedRunner: failed to create TensorRT runtime for %s", enginePath.c_str()); + return false; + } + + mEngine.reset(mRuntime->deserializeCudaEngine(engineFileReader.getData(), engineFileReader.getSize())); + if (!mEngine) + { + LOG_ERROR("EmbedRunner: failed to deserialize TensorRT engine: %s", enginePath.c_str()); + return false; + } + + mContext.reset(mEngine->createExecutionContext(nvinfer1::ExecutionContextAllocationStrategy::kUSER_MANAGED)); + if (!mContext) + { + LOG_ERROR("EmbedRunner: failed to create TensorRT execution context for %s", enginePath.c_str()); + return false; + } + + if (!mContext->setOptimizationProfileAsync(0, stream)) + { + LOG_ERROR("EmbedRunner: failed to set optimization profile for %s", enginePath.c_str()); + return false; + } + CUDA_CHECK(cudaStreamSynchronize(stream)); + return true; +} + +bool EmbedRunner::validateAndFillConfig() +{ + auto const modelType = mConfigJson.value("model_type", std::string{"cosmos_embed"}); + auto const component = mConfigJson.value("component", std::string{"embed"}); + if (modelType != "cosmos_embed" && component != "embed") + { + LOG_ERROR("EmbedRunner: unexpected model_type=%s component=%s", modelType.c_str(), component.c_str()); + return false; + } + + if (mConfigJson.contains("input_names") && mConfigJson.at("input_names").size() >= 2U) + { + mLatentsInputName = mConfigJson.at("input_names").at(0).get(); + mTimestepInputName = mConfigJson.at("input_names").at(1).get(); + } + if (mConfigJson.contains("output_names") && !mConfigJson.at("output_names").empty()) + { + mOutputName = mConfigJson.at("output_names").at(0).get(); + } + + if (!mConfigJson.contains("inputs")) + { + LOG_ERROR("EmbedRunner: config is missing inputs metadata"); + return false; + } + if (!mConfigJson.contains("outputs") || mConfigJson.at("outputs").empty()) + { + LOG_ERROR("EmbedRunner: config is missing output metadata"); + return false; + } + + auto const& inputs = mConfigJson.at("inputs"); + if (!inputs.contains(mLatentsInputName) || !inputs.contains(mTimestepInputName)) + { + LOG_ERROR("EmbedRunner: config is missing input metadata for %s and/or %s", mLatentsInputName.c_str(), + mTimestepInputName.c_str()); + return false; + } + + auto const& latentsMeta = inputs.at(mLatentsInputName); + auto const& timestepMeta = inputs.at(mTimestepInputName); + auto const& outputMeta = mConfigJson.at("outputs").at(0); + + mLatentsShape = coordsFromJson(latentsMeta.at("shape")); + if (timestepMeta.at("shape").empty()) + { + mTimestepIsScalar = true; + mTimestepShape = rt::Coords({1}); + } + else + { + mTimestepShape = coordsFromJson(timestepMeta.at("shape")); + } + mOutputShape = coordsFromJson(outputMeta.at("shape")); + mLatentsType = dataTypeFromTorchString(latentsMeta.at("dtype").get()); + mTimestepType = dataTypeFromTorchString(timestepMeta.at("dtype").get()); + mOutputType = dataTypeFromTorchString(outputMeta.at("dtype").get()); + + if (mLatentsShape.getNumDims() != 5) + { + LOG_ERROR("EmbedRunner: expected rank-5 vision latents, got %d dims", mLatentsShape.getNumDims()); + return false; + } + if (mOutputShape.getNumDims() != 2) + { + LOG_ERROR("EmbedRunner: expected rank-2 gen_seq output, got %d dims", mOutputShape.getNumDims()); + return false; + } + + mLatentsInputName = resolveIOTensorName(mEngine.get(), mLatentsInputName, nvinfer1::TensorIOMode::kINPUT, 0); + mTimestepInputName = resolveIOTensorName(mEngine.get(), mTimestepInputName, nvinfer1::TensorIOMode::kINPUT, 1); + mOutputName = resolveIOTensorName(mEngine.get(), mOutputName, nvinfer1::TensorIOMode::kOUTPUT, 0); + return true; +} + +int64_t EmbedRunner::getRequiredContextMemorySize() const +{ + return mEngine ? mEngine->getDeviceMemorySizeV2() : 0; +} + +bool EmbedRunner::allocateBuffers() +{ + mLatentsTensor = rt::Tensor(mLatentsShape, rt::DeviceType::kGPU, mLatentsType, mLatentsInputName); + mTimestepTensor = rt::Tensor(mTimestepShape, rt::DeviceType::kGPU, mTimestepType, mTimestepInputName); + mGenSeqTensor = rt::Tensor(mOutputShape, rt::DeviceType::kGPU, mOutputType, mOutputName); + return bindTensors(); +} + +bool EmbedRunner::bindTensors() noexcept +{ + if (!mContext->setInputShape(mLatentsInputName.c_str(), mLatentsTensor.getShape().getTRTDims())) + { + LOG_ERROR("EmbedRunner: failed to set input shape for %s", mLatentsInputName.c_str()); + return false; + } + nvinfer1::Dims timestepDims = mTimestepIsScalar ? nvinfer1::Dims{} : mTimestepTensor.getShape().getTRTDims(); + if (!mContext->setInputShape(mTimestepInputName.c_str(), timestepDims)) + { + LOG_ERROR("EmbedRunner: failed to set input shape for %s", mTimestepInputName.c_str()); + return false; + } + if (!mContext->setTensorAddress(mLatentsInputName.c_str(), mLatentsTensor.rawPointer())) + { + LOG_ERROR("EmbedRunner: failed to bind input tensor %s", mLatentsInputName.c_str()); + return false; + } + if (!mContext->setTensorAddress(mTimestepInputName.c_str(), mTimestepTensor.rawPointer())) + { + LOG_ERROR("EmbedRunner: failed to bind input tensor %s", mTimestepInputName.c_str()); + return false; + } + if (!mContext->setTensorAddress(mOutputName.c_str(), mGenSeqTensor.rawPointer())) + { + LOG_ERROR("EmbedRunner: failed to bind output tensor %s", mOutputName.c_str()); + return false; + } + return true; +} + +bool EmbedRunner::setContextMemory(rt::Tensor& sharedContextMemory) +{ + if (!mEngine) + { + return true; + } + + int64_t const requiredSize = getRequiredContextMemorySize(); + if (sharedContextMemory.getMemoryCapacity() < requiredSize) + { + LOG_ERROR("EmbedRunner: shared context memory (%zu bytes) is smaller than required (%zu bytes)", + static_cast(sharedContextMemory.getMemoryCapacity()), static_cast(requiredSize)); + return false; + } + + mContext->setDeviceMemoryV2(sharedContextMemory.rawPointer(), sharedContextMemory.getMemoryCapacity()); + return true; +} + +bool EmbedRunner::copyLatentsFrom(rt::Tensor const& src, cudaStream_t stream) +{ + if (src.getShape() != mLatentsTensor.getShape()) + { + LOG_ERROR("EmbedRunner: latent shape %s does not match engine input %s", src.getShape().formatString().c_str(), + mLatentsTensor.getShape().formatString().c_str()); + return false; + } + return copyTensorToDevice(mLatentsTensor, src, stream); +} + +bool EmbedRunner::setTimestep(float timestep, cudaStream_t stream) +{ + if (mTimestepType == nvinfer1::DataType::kFLOAT) + { + if (mTimestepShape.volume() == 1) + { + CUDA_CHECK(cudaMemcpyAsync( + mTimestepTensor.rawPointer(), ×tep, sizeof(float), cudaMemcpyHostToDevice, stream)); + return true; + } + std::vector values(static_cast(mTimestepShape.volume()), timestep); + CUDA_CHECK(cudaMemcpyAsync(mTimestepTensor.rawPointer(), values.data(), values.size() * sizeof(float), + cudaMemcpyHostToDevice, stream)); + return true; + } + + if (mTimestepType == nvinfer1::DataType::kHALF) + { + half const value = __float2half(timestep); + if (mTimestepShape.volume() == 1) + { + CUDA_CHECK( + cudaMemcpyAsync(mTimestepTensor.rawPointer(), &value, sizeof(half), cudaMemcpyHostToDevice, stream)); + return true; + } + std::vector values(static_cast(mTimestepShape.volume()), value); + CUDA_CHECK(cudaMemcpyAsync( + mTimestepTensor.rawPointer(), values.data(), values.size() * sizeof(half), cudaMemcpyHostToDevice, stream)); + return true; + } + + LOG_ERROR("EmbedRunner: unsupported timestep dtype."); + return false; +} + +bool EmbedRunner::embed(cudaStream_t stream) noexcept +{ + if (!bindTensors()) + { + return false; + } + if (!mContext->enqueueV3(stream)) + { + LOG_ERROR("EmbedRunner: enqueueV3 failed."); + return false; + } + return true; +} + +} // namespace rt +} // namespace trt_edgellm diff --git a/cpp/runtime/embedRunner.h b/cpp/runtime/embedRunner.h new file mode 100644 index 00000000..0b041441 --- /dev/null +++ b/cpp/runtime/embedRunner.h @@ -0,0 +1,143 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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. + */ + +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "common/tensor.h" +#include "runtime/wfmRuntimeUtils.h" + +#include +#include +#include + +#include +#include +#include + +namespace trt_edgellm +{ +namespace rt +{ + +//! Cosmos vision embed engine: latents + timestep -> gen_seq vision tokens. +class EmbedRunner +{ +public: + //! \p engineDir is the ``embed/`` component directory (contains config.json + engine). + explicit EmbedRunner(std::string const& engineDir, cudaStream_t stream); + + ~EmbedRunner() noexcept = default; + + int64_t getRequiredContextMemorySize() const; + bool setContextMemory(rt::Tensor& sharedContextMemory); + + std::string const& getLatentsInputName() const noexcept + { + return mLatentsInputName; + } + + std::string const& getTimestepInputName() const noexcept + { + return mTimestepInputName; + } + + std::string const& getOutputName() const noexcept + { + return mOutputName; + } + + rt::Coords const& getLatentShape() const noexcept + { + return mLatentsShape; + } + + rt::Coords const& getGenSeqShape() const noexcept + { + return mOutputShape; + } + + rt::Tensor& getLatents() noexcept + { + return mLatentsTensor; + } + + rt::Tensor const& getLatents() const noexcept + { + return mLatentsTensor; + } + + rt::Tensor& getTimestep() noexcept + { + return mTimestepTensor; + } + + rt::Tensor const& getTimestep() const noexcept + { + return mTimestepTensor; + } + + rt::Tensor& getGenSeq() noexcept + { + return mGenSeqTensor; + } + + rt::Tensor const& getGenSeq() const noexcept + { + return mGenSeqTensor; + } + + bool copyLatentsFrom(rt::Tensor const& src, cudaStream_t stream); + bool setTimestep(float timestep, cudaStream_t stream); + bool embed(cudaStream_t stream) noexcept; + +private: + bool loadConfig(); + bool loadEngine(cudaStream_t stream); + bool validateAndFillConfig(); + bool allocateBuffers(); + bool bindTensors() noexcept; + + std::string mEngineDir; + nlohmann::json mConfigJson; + + std::unique_ptr mRuntime; + std::unique_ptr mEngine; + std::unique_ptr mContext; + + std::string mLatentsInputName{"vision_latents"}; + std::string mTimestepInputName{"timestep"}; + std::string mOutputName{"gen_seq"}; + rt::Coords mLatentsShape; + rt::Coords mTimestepShape; + rt::Coords mOutputShape; + bool mTimestepIsScalar{false}; + nvinfer1::DataType mLatentsType{nvinfer1::DataType::kHALF}; + nvinfer1::DataType mTimestepType{nvinfer1::DataType::kFLOAT}; + nvinfer1::DataType mOutputType{nvinfer1::DataType::kHALF}; + + rt::Tensor mLatentsTensor; + rt::Tensor mTimestepTensor; + rt::Tensor mGenSeqTensor; +}; + +} // namespace rt +} // namespace trt_edgellm diff --git a/cpp/runtime/linearKVCache.cpp b/cpp/runtime/linearKVCache.cpp new file mode 100644 index 00000000..4c9f8062 --- /dev/null +++ b/cpp/runtime/linearKVCache.cpp @@ -0,0 +1,324 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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 "runtime/linearKVCache.h" +#include "common/logger.h" + +#include "common/checkMacros.h" +#include "common/cudaMacros.h" +#include "kernels/kvCacheUtilKernels/kvCacheUtilsKernels.h" +#include + +using namespace nvinfer1; + +namespace trt_edgellm +{ +namespace rt +{ + +LinearKVCache::LinearKVCache(CacheConfig const& config, cudaStream_t stream) + : mConfig(config) +{ + check::check( + mConfig.kvCacheTypeTRT == nvinfer1::DataType::kHALF || mConfig.kvCacheTypeTRT == nvinfer1::DataType::kFP8, + "Unsupported KV cache dtype."); + mDeviceKVCache = rt::Tensor({mConfig.numAttentionLayers, mConfig.maxBatchSize, 2, mConfig.numKVHeads, + mConfig.maxSequenceLength, mConfig.headDim}, + DeviceType::kGPU, mConfig.kvCacheTypeTRT, "LinearKVCache::mDeviceKVCache"); + int64_t const kvCacheVolume = mConfig.numAttentionLayers * mConfig.maxBatchSize * 2 * mConfig.numKVHeads + * mConfig.maxSequenceLength * mConfig.headDim; + size_t const kvCacheElemSize = rt::utils::getTypeSize(mConfig.kvCacheTypeTRT); + char const* kvCacheTypeStr = (mConfig.kvCacheTypeTRT == nvinfer1::DataType::kHALF) ? "kHALF" : "kFP8"; + LOG_DEBUG( + "KVCache(dtype=%s) of shape [%ld, %ld, %ld, %ld, %ld, %ld] allocated on GPU with size: %ld bytes (%.2f MB)", + kvCacheTypeStr, mConfig.numAttentionLayers, mConfig.maxBatchSize, 2, mConfig.numKVHeads, + mConfig.maxSequenceLength, mConfig.headDim, kvCacheVolume * static_cast(kvCacheElemSize), + static_cast(kvCacheVolume * static_cast(kvCacheElemSize)) / (1024.0 * 1024.0)); + mDeviceKVCacheLengths = rt::Tensor( + {mConfig.maxBatchSize}, DeviceType::kGPU, DataType::kINT32, "LinearKVCache::mDeviceKVCacheLengths"); + CUDA_CHECK( + cudaMemsetAsync(mDeviceKVCacheLengths.rawPointer(), 0, mDeviceKVCacheLengths.getMemoryCapacity(), stream)); + + if (mConfig.numLinearAttnLayers > 0) + { + mDeviceRecurrentStates + = rt::Tensor({mConfig.numLinearAttnLayers, mConfig.maxBatchSize, mConfig.recurrentStateNumHeads, + mConfig.recurrentStateHeadDim, mConfig.recurrentStateSize}, + DeviceType::kGPU, mConfig.recurrentStateType, "LinearKVCache::mDeviceRecurrentStates"); + CUDA_CHECK(cudaMemsetAsync( + mDeviceRecurrentStates.rawPointer(), 0, mDeviceRecurrentStates.getMemoryCapacity(), stream)); + + mDeviceConvStates + = rt::Tensor({mConfig.numLinearAttnLayers, mConfig.maxBatchSize, mConfig.convDim, mConfig.convKernel}, + DeviceType::kGPU, mConfig.convStateType, "LinearKVCache::mDeviceConvStates"); + CUDA_CHECK(cudaMemsetAsync(mDeviceConvStates.rawPointer(), 0, mDeviceConvStates.getMemoryCapacity(), stream)); + } +} + +LinearKVCache::~LinearKVCache() noexcept {} + +LinearKVCache::LinearKVCache(LinearKVCache&& other) noexcept +{ + mConfig = other.mConfig; + mActiveBatchSize = other.mActiveBatchSize; + mKVCacheAllEmpty = other.mKVCacheAllEmpty; + mDeviceKVCache = std::move(other.mDeviceKVCache); + mDeviceKVCacheLengths = std::move(other.mDeviceKVCacheLengths); + mDeviceRecurrentStates = std::move(other.mDeviceRecurrentStates); + mDeviceConvStates = std::move(other.mDeviceConvStates); + + other.mConfig = CacheConfig{}; + other.mActiveBatchSize = 0; + other.mKVCacheAllEmpty = true; +} + +LinearKVCache& LinearKVCache::operator=(LinearKVCache&& other) noexcept +{ + if (this != &other) + { + mConfig = other.mConfig; + mKVCacheAllEmpty = other.mKVCacheAllEmpty; + mActiveBatchSize = other.mActiveBatchSize; + mDeviceKVCache = std::move(other.mDeviceKVCache); + mDeviceKVCacheLengths = std::move(other.mDeviceKVCacheLengths); + mDeviceRecurrentStates = std::move(other.mDeviceRecurrentStates); + mDeviceConvStates = std::move(other.mDeviceConvStates); + + other.mConfig = CacheConfig{}; + other.mActiveBatchSize = 0; + other.mKVCacheAllEmpty = true; + } + return *this; +} + +rt::Tensor LinearKVCache::getCombinedKVCacheForDecoderLayer(int32_t decoderLayerIdx) noexcept +{ + int64_t const kvCacheOffset + = decoderLayerIdx * mConfig.maxBatchSize * 2 * mConfig.numKVHeads * mConfig.maxSequenceLength * mConfig.headDim; + + size_t const elemSize = rt::utils::getTypeSize(mConfig.kvCacheTypeTRT); + void* kvCachePtr = static_cast(static_cast(mDeviceKVCache.rawPointer()) + kvCacheOffset * elemSize); + + return rt::Tensor(kvCachePtr, + {mConfig.maxBatchSize, 2, mConfig.numKVHeads, mConfig.maxSequenceLength, mConfig.headDim}, DeviceType::kGPU, + mConfig.kvCacheTypeTRT); +} + +std::pair LinearKVCache::getSeparateKVCacheForDecoderLayer(int32_t decoderLayerIdx) noexcept +{ + // Get the combined KV cache for this layer from base class + rt::Tensor kvCache = LinearKVCache::getCombinedKVCacheForDecoderLayer(decoderLayerIdx); + + // The KV cache has shape: [maxBatchSize, 2, numKVHeads, maxSequenceLength, headDim] + // K cache is at index 0 of dimension 2, so we just need to point to the beginning + // and reshape to remove the "2" dimension + + CacheConfig config = getConfig(); + void* kvCachePtr = static_cast(kvCache.rawPointer()); + + void* kCachePtr = static_cast(kvCachePtr); + rt::Tensor kCache + = rt::Tensor(kCachePtr, {config.maxBatchSize, config.numKVHeads, config.maxSequenceLength, config.headDim}, + DeviceType::kGPU, mConfig.kvCacheTypeTRT); + + // Calculate offset to V cache: skip the entire K cache portion + // Offset = maxBatchSize * 1 (K portion) * numKVHeads * maxSequenceLength * headDim + int64_t vCacheOffset = config.maxBatchSize * config.numKVHeads * config.maxSequenceLength * config.headDim + * static_cast(rt::utils::getTypeSize(mConfig.kvCacheTypeTRT)); + void* vCachePtr = static_cast(static_cast(kvCachePtr) + vCacheOffset); + + rt::Tensor vCache + = rt::Tensor(vCachePtr, {config.maxBatchSize, config.numKVHeads, config.maxSequenceLength, config.headDim}, + DeviceType::kGPU, mConfig.kvCacheTypeTRT); + return {std::move(kCache), std::move(vCache)}; +} + +rt::Tensor LinearKVCache::getKVCacheBuffer() noexcept +{ + return rt::Tensor(mDeviceKVCache.rawPointer(), + {mConfig.numAttentionLayers, mConfig.maxBatchSize, 2, mConfig.numKVHeads, mConfig.maxSequenceLength, + mConfig.headDim}, + DeviceType::kGPU, mConfig.kvCacheTypeTRT); +} + +void LinearKVCache::clearRecurrentStates(cudaStream_t stream) +{ + if (mConfig.numLinearAttnLayers == 0) + { + return; + } + CUDA_CHECK( + cudaMemsetAsync(mDeviceRecurrentStates.rawPointer(), 0, mDeviceRecurrentStates.getMemoryCapacity(), stream)); + CUDA_CHECK(cudaMemsetAsync(mDeviceConvStates.rawPointer(), 0, mDeviceConvStates.getMemoryCapacity(), stream)); +} + +rt::Tensor LinearKVCache::getRecurrentStateForLayer(int32_t recurrentLayerIdx) noexcept +{ + size_t const elemSize = rt::utils::getTypeSize(mConfig.recurrentStateType); + int64_t const perLayerElems = mConfig.maxBatchSize * mConfig.recurrentStateNumHeads * mConfig.recurrentStateHeadDim + * mConfig.recurrentStateSize; + void* ptr = static_cast(mDeviceRecurrentStates.rawPointer()) + recurrentLayerIdx * perLayerElems * elemSize; + return rt::Tensor(ptr, + {mConfig.maxBatchSize, mConfig.recurrentStateNumHeads, mConfig.recurrentStateHeadDim, + mConfig.recurrentStateSize}, + DeviceType::kGPU, mConfig.recurrentStateType); +} + +rt::Tensor LinearKVCache::getConvStateForLayer(int32_t recurrentLayerIdx) noexcept +{ + size_t const elemSize = rt::utils::getTypeSize(mConfig.convStateType); + int64_t const perLayerElems = mConfig.maxBatchSize * mConfig.convDim * mConfig.convKernel; + void* ptr = static_cast(mDeviceConvStates.rawPointer()) + recurrentLayerIdx * perLayerElems * elemSize; + return rt::Tensor( + ptr, {mConfig.maxBatchSize, mConfig.convDim, mConfig.convKernel}, DeviceType::kGPU, mConfig.convStateType); +} + +std::vector LinearKVCache::captureRecurrentStates(int32_t batchIdx, cudaStream_t stream) +{ + std::vector result; + if (mConfig.numLinearAttnLayers == 0) + { + return result; + } + size_t const elemSize = rt::utils::getTypeSize(mConfig.recurrentStateType); + int64_t const perLayerElems = mConfig.maxBatchSize * mConfig.recurrentStateNumHeads * mConfig.recurrentStateHeadDim + * mConfig.recurrentStateSize; + int64_t const perBatchElems + = mConfig.recurrentStateNumHeads * mConfig.recurrentStateHeadDim * mConfig.recurrentStateSize; + size_t const perBatchBytes = static_cast(perBatchElems) * elemSize; + + result.reserve(mConfig.numLinearAttnLayers); + for (int32_t layer = 0; layer < mConfig.numLinearAttnLayers; ++layer) + { + void const* src = static_cast(mDeviceRecurrentStates.rawPointer()) + + static_cast(layer * perLayerElems + batchIdx * perBatchElems) * elemSize; + rt::Tensor saved({1, mConfig.recurrentStateNumHeads, mConfig.recurrentStateHeadDim, mConfig.recurrentStateSize}, + DeviceType::kGPU, mConfig.recurrentStateType, + "LinearKVCache::capturedRecurrentState_" + std::to_string(layer)); + CUDA_CHECK(cudaMemcpyAsync(saved.rawPointer(), src, perBatchBytes, cudaMemcpyDeviceToDevice, stream)); + result.push_back(std::move(saved)); + } + return result; +} + +std::vector LinearKVCache::captureConvStates(int32_t batchIdx, cudaStream_t stream) +{ + std::vector result; + if (mConfig.numLinearAttnLayers == 0) + { + return result; + } + size_t const elemSize = rt::utils::getTypeSize(mConfig.convStateType); + int64_t const perLayerElems = mConfig.maxBatchSize * mConfig.convDim * mConfig.convKernel; + int64_t const perBatchElems = mConfig.convDim * mConfig.convKernel; + size_t const perBatchBytes = static_cast(perBatchElems) * elemSize; + + result.reserve(mConfig.numLinearAttnLayers); + for (int32_t layer = 0; layer < mConfig.numLinearAttnLayers; ++layer) + { + void const* src = static_cast(mDeviceConvStates.rawPointer()) + + static_cast(layer * perLayerElems + batchIdx * perBatchElems) * elemSize; + rt::Tensor saved({1, mConfig.convDim, mConfig.convKernel}, DeviceType::kGPU, mConfig.convStateType, + "LinearKVCache::capturedConvState_" + std::to_string(layer)); + CUDA_CHECK(cudaMemcpyAsync(saved.rawPointer(), src, perBatchBytes, cudaMemcpyDeviceToDevice, stream)); + result.push_back(std::move(saved)); + } + return result; +} + +void LinearKVCache::resetForNewSequences(rt::Tensor const& reuseKVCacheLengths, cudaStream_t stream) +{ + int32_t const batchSize = static_cast(reuseKVCacheLengths.getShape()[0]); + check::check( + batchSize <= mConfig.maxBatchSize, "Batch size of request shall not exceed the max supported batch size."); + check::check( + reuseKVCacheLengths.getDeviceType() == DeviceType::kCPU, "The reuseKVCacheLengths tensor shall reside on CPU."); + check::check(reuseKVCacheLengths.getDataType() == mDeviceKVCacheLengths.getDataType(), + "The data type of the reuseKVCacheLengths tensor shall match the data type of the Device KVCache Lengths."); + + mActiveBatchSize = batchSize; + check::check(mDeviceKVCacheLengths.reshape({mActiveBatchSize}), "Tensor reshape failed"); + + // If all reuseSequenceLengths are 0, then we can set flag mKVCacheAllEmpty to true. + int32_t const* reuseSequenceLengthsData = reuseKVCacheLengths.dataPointer(); + bool allEmpty{true}; + for (int32_t i = 0; i < batchSize; ++i) + { + if (reuseSequenceLengthsData[i] != 0) + { + allEmpty = false; + break; + } + } + mKVCacheAllEmpty = allEmpty; + CUDA_CHECK(cudaMemcpyAsync(mDeviceKVCacheLengths.rawPointer(), reuseKVCacheLengths.rawPointer(), + reuseKVCacheLengths.getMemoryCapacity(), cudaMemcpyHostToDevice, stream)); +} + +void LinearKVCache::commitSequenceLength(rt::Tensor const& newContextLengths, cudaStream_t stream) +{ + check::check(newContextLengths.getDataType() == DataType::kINT32, + "The newContextLengths tensor shall have data type of int32_t."); + check::check( + newContextLengths.getDeviceType() == DeviceType::kGPU, "The newContextLengths tensor shall reside on GPU."); + check::check(newContextLengths.getShape()[0] == mActiveBatchSize, + "The newContextLengths tensor shall have the same batch size as the active batch size."); + + kernel::incrementLengthTensor(mDeviceKVCacheLengths, newContextLengths, stream); + + // Set flag to false since we have committed a new sequence length. + mKVCacheAllEmpty = false; +} + +void LinearKVCache::commitSequenceLength(int32_t increment, cudaStream_t stream) +{ + kernel::incrementLengthTensor(mDeviceKVCacheLengths, increment, stream); + + // Set flag to false since we have committed a new sequence length. + mKVCacheAllEmpty = false; +} + +rt::Tensor& LinearKVCache::getKVCacheLengths() noexcept +{ + return mDeviceKVCacheLengths; +} + +LinearKVCache::CacheConfig LinearKVCache::getConfig() const noexcept +{ + return mConfig; +} + +int32_t LinearKVCache::getActiveBatchSize() const noexcept +{ + return mActiveBatchSize; +} + +bool LinearKVCache::getKVCacheAllEmpty() const noexcept +{ + return mKVCacheAllEmpty; +} + +void LinearKVCache::setActiveBatchSize(int32_t newActiveBatchSize) +{ + check::check(newActiveBatchSize >= 0 && newActiveBatchSize <= mConfig.maxBatchSize, + "Invalid active batch size: must be in range [0, maxBatchSize]"); + mActiveBatchSize = newActiveBatchSize; + check::check(mDeviceKVCacheLengths.reshape({mActiveBatchSize}), "Tensor reshape failed"); +} + +} // namespace rt +} // namespace trt_edgellm diff --git a/cpp/runtime/linearKVCache.h b/cpp/runtime/linearKVCache.h new file mode 100644 index 00000000..ed5b1be6 --- /dev/null +++ b/cpp/runtime/linearKVCache.h @@ -0,0 +1,187 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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. + */ + +#pragma once + +#include +#include +#include +namespace trt_edgellm +{ +namespace rt +{ + +//! Static Linear KVCache that holds the KVCache for all decoder layers up to maxSequenceLength. +//! The KVCache implement the design of: +//! 1. Allocates memory for max supported batch size. +//! 2. Memory Layout: [numAttentionLayers, maxBatchSize, 2, numKVHeads, maxSequenceLength, headDim] +//! 3. Synchronous execution of batch requests, all the sequences in the batch will run prefill +//! or decode at the same time. +class LinearKVCache +{ +public: + //! \cond INTERNAL + /*! + * @brief Configuration for KV cache + * + * Defines the dimensions and capacity of the KV cache. + */ + struct CacheConfig + { + int64_t numAttentionLayers{}; //!< Number of attention layers needing KV cache + int64_t maxBatchSize{}; //!< Maximum batch size + int64_t maxSequenceLength{}; //!< Maximum sequence length + int64_t numKVHeads{}; //!< Number of key-value heads + int64_t headDim{}; //!< Head dimension + nvinfer1::DataType kvCacheTypeTRT{}; //!< Storage dtype for KV cache (kHALF or kFP8) + + // Recurrent state fields (all zero for pure-attention models; no memory is allocated) + int32_t numLinearAttnLayers{0}; //!< Number of recurrent layers + int32_t recurrentStateNumHeads{0}; //!< Number of recurrent state heads + int32_t recurrentStateHeadDim{0}; //!< Dimension of each recurrent head + int32_t recurrentStateSize{0}; //!< Recurrent state dimension + nvinfer1::DataType recurrentStateType{nvinfer1::DataType::kHALF}; //!< Recurrent state dtype + int32_t convDim{0}; //!< Conv1d channel dimension + int32_t convKernel{0}; //!< Conv1d kernel width + nvinfer1::DataType convStateType{nvinfer1::DataType::kHALF}; //!< Conv state dtype + }; + //! \endcond + + //! @brief Default constructor + LinearKVCache() noexcept = default; + + /*! + * @brief Construct and initialize KV cache + * + * Allocates device memory for KV cache. Once allocated, memory won't be reallocated. + * + * @param config Cache configuration + * @param stream CUDA stream for allocation + * @throws std::runtime_error if CUDA operations fail or data type is unsupported + */ + LinearKVCache(CacheConfig const& config, cudaStream_t stream); + + //! @brief Destructor + ~LinearKVCache() noexcept; + + //! @brief Deleted copy constructor to avoid large data copy + LinearKVCache(LinearKVCache const&) = delete; + + //! @brief Deleted copy assignment to avoid large data copy + //! @return Reference to this + LinearKVCache& operator=(LinearKVCache const&) = delete; + + //! @brief Move constructor + LinearKVCache(LinearKVCache&&) noexcept; + + //! @brief Move assignment operator + //! @return Reference to this + LinearKVCache& operator=(LinearKVCache&&) noexcept; + + //! Get the combined KVCache for the given decoder layer, for EdgeLLM Attention TRT plugin implementation. + //! @param decoderLayerIdx The index of the decoder layer. + //! @return A non-owned tensor object with shape [batch_size, 2, num_kv_heads, max_sequence_length, head_dim] that + //! points to the combined KVCache memory with shape information. + rt::Tensor getCombinedKVCacheForDecoderLayer(int32_t decoderLayerIdx) noexcept; + + //! Get the separate K and V caches for the given decoder layer, for TRT native KVCacheUpdate/Attention operations. + //! Returns a pair of tensors, the first is the K cache and the second is the V cache. + //! @param decoderLayerIdx The index of the decoder layer. + //! @return A pair of tensors, the first is the K cache and the second is the V cache, with shapes [batch_size, + //! num_kv_heads, max_sequence_length, head_dim]. + std::pair getSeparateKVCacheForDecoderLayer(int32_t decoderLayerIdx) noexcept; + + //! Get the full KVCache buffer as a non-owned tensor. + rt::Tensor getKVCacheBuffer() noexcept; + + //! Get recurrent state tensor for a recurrent layer (non-owned view). + //! Shape: [maxBatchSize, recurrentStateNumHeads, recurrentStateHeadDim, recurrentStateSize] + rt::Tensor getRecurrentStateForLayer(int32_t recurrentLayerIdx) noexcept; + + //! Get conv state tensor for a recurrent layer (non-owned view). + //! Shape: [maxBatchSize, convDim, convKernel] + rt::Tensor getConvStateForLayer(int32_t recurrentLayerIdx) noexcept; + + //! Zero all recurrent and conv state buffers (all layers, all batch slots). + //! Called after warmup inference and before CUDA graph capture to ensure a clean starting state. + void clearRecurrentStates(cudaStream_t stream); + + //! Copy one batch slot's recurrent states into freshly-allocated tensors (one per recurrent layer). + //! Used to snapshot states when saving a system prompt cache entry. + std::vector captureRecurrentStates(int32_t batchIdx, cudaStream_t stream); + + //! Copy one batch slot's conv states into freshly-allocated tensors (one per recurrent layer). + //! Used to snapshot states when saving a system prompt cache entry. + std::vector captureConvStates(int32_t batchIdx, cudaStream_t stream); + + //! Asynchronously reset the KVCache buffer state for a new setup of input context. + //! @param hostReuseKVCacheLengths The lengths of the KVCache to be reused from precomputed KVCache content. + //! @param stream The stream is used to perform GPU memory operations. + //! @throws std::runtime_error if tensor shape, location or data type are invalid, or if a CUDA operation fails + void resetForNewSequences(rt::Tensor const& hostReuseKVCacheLengths, cudaStream_t stream); + + //! Asynchronously commit the KVCache buffer for a prefill request, record stored KVCache lengths. + //! @param newContextLengths [GPU, Int32]: The context length to commit for the KVCache. + //! @param stream The stream is used to perform GPU memory operations. + //! @throws std::runtime_error if tensor shape, location or data type are invalid + void commitSequenceLength(rt::Tensor const& newContextLengths, cudaStream_t stream); + + //! Commit the KVCache buffer for a decode request, increment the KVCache lengths by 1 for active sequences. + //! @param increment The amount to increment sequence lengths (typically 1 for decode step) + //! @param stream The stream is used to perform GPU memory operations. + //! @throws std::runtime_error if KV cache lengths tensor has wrong location or data type + void commitSequenceLength(int32_t increment, cudaStream_t stream); + + //! @brief Get KV cache lengths for active sequences + //! @return Reference to KV cache lengths tensor + rt::Tensor& getKVCacheLengths() noexcept; + + //! @brief Get KV cache configuration + //! @return Cache configuration + CacheConfig getConfig() const noexcept; + + //! @brief Get active batch size + //! @return Number of active sequences + int32_t getActiveBatchSize() const noexcept; + + //! @brief Get flag to indicate if KVCache for all sequences are empty. + //! @return Flag to indicate if KVCache for all sequences are empty. + bool getKVCacheAllEmpty() const noexcept; + + //! @brief Set active batch size (for batch eviction) + //! @param newActiveBatchSize New active batch size after eviction + //! @throws std::runtime_error If newActiveBatchSize is out of valid range [0, maxBatchSize] + void setActiveBatchSize(int32_t newActiveBatchSize); + +private: + CacheConfig mConfig{}; //!< Cache configuration + int32_t mActiveBatchSize{}; //!< Active batch size + bool mKVCacheAllEmpty{}; //!< Flag to indicate if KVCache for all sequences are empty. + rt::Tensor mDeviceKVCacheLengths{}; //!< KV cache lengths on device + rt::Tensor mDeviceKVCache{}; //!< KV cache memory buffer on device + + //! Recurrent state buffer: [numLinearAttnLayers, maxBatchSize, recurrentStateNumHeads, recurrentStateHeadDim, + //! recurrentStateSize] Empty when numLinearAttnLayers == 0. + rt::Tensor mDeviceRecurrentStates{}; + + //! Conv state buffer: [numLinearAttnLayers, maxBatchSize, convDim, convKernel] + //! Empty when numLinearAttnLayers == 0. + rt::Tensor mDeviceConvStates{}; +}; + +} // namespace rt +} // namespace trt_edgellm diff --git a/cpp/runtime/llmEngineRunner.cpp b/cpp/runtime/llmEngineRunner.cpp new file mode 100644 index 00000000..b65807e3 --- /dev/null +++ b/cpp/runtime/llmEngineRunner.cpp @@ -0,0 +1,2362 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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 "runtime/llmEngineRunner.h" + +#include "common/bindingNames.h" +#include "common/checkMacros.h" +#include "common/cudaUtils.h" +#include "common/hashUtils.h" +#include "common/logger.h" +#include "common/mmapReader.h" +#include "common/safetensorsUtils.h" +#include "common/stringUtils.h" +#include "common/trtUtils.h" +#include "common/version.h" +#include "kernels/embeddingKernels/embeddingKernels.h" +#include "kernels/kvCacheUtilKernels/kvCacheUtilsKernels.h" +#include "kernels/posEncoding/initializeCosSinCache.h" +#include "kernels/speculative/eagleUtilKernels.h" +#include "profiling/layerProfiler.h" +#include "runtime/llmRuntimeUtils.h" +#include +#include +#include + +using namespace trt_edgellm; +using namespace nvinfer1; + +namespace +{ +//! Dummy dimension for LoRA weights when no LoRA is active (use 1 instead of 0 to avoid zero-shape issues) +constexpr int32_t kEMPTY_LORA_RANK = 1; + +std::string formatEngineConfig(trt_edgellm::rt::LLMEngineRunnerConfig const& config) +{ + std::stringstream ss; + + ss << std::boolalpha; + ss << "LLMEngineRunnerConfig:" + << " enableEagleSpecDecode: " << config.enableEagleSpecDecode + << " numDecoderLayers: " << config.numDecoderLayers << " numKVHeads: " << config.numKVHeads + << " headDim: " << config.headDim << " rotaryDim: " << config.rotaryDim + << " hiddenSize: " << config.hiddenSize << " maxSupportedBatchSize: " << config.maxSupportedBatchSize + << " maxSupportedInputLength: " << config.maxSupportedInputLength + << " maxKVCacheCapacity: " << config.maxKVCacheCapacity + << " maxSupportedLoraRank: " << config.maxSupportedLoraRank + << " numDeepstackFeatures: " << config.numDeepstackFeatures; + if (config.enableEagleSpecDecode) + { + ss << " outputHiddenDim (For Eagle SpecDecode): " << config.outputHiddenDim; + ss << " maxVerifyTreeSize (For Eagle SpecDecode): " << config.maxVerifyTreeSize; + } + if (config.enableContextEmb) + { + ss << " contextEmbDim (For context_emb): " << config.contextEmbDim; + } + if (config.enableLmHiddenStates) + { + ss << " enableLmHiddenStates: true"; + } + return ss.str(); +} + +// Compute a unique key value that can distinguish the various decoding steps. +// Extend this function when we need to capture more information. +trt_edgellm::rt::LLMEngineRunner::DecodingGraphKey decodingKey( + rt::Tensor const& inputsEmbeds, rt::Tensor const& outputLogits, std::string const& loraWeightsName) noexcept +{ + // For vanilla decoding step, the shape can be distingusihed by active batch size. + // Also capture the pointer address to ensure we are read/write correct locations. + int64_t const activeBatchSize = inputsEmbeds.getShape()[0]; + uintptr_t const inputsEmbedsAddr = reinterpret_cast(inputsEmbeds.rawPointer()); + uintptr_t const outputLogitsAddr = reinterpret_cast(outputLogits.rawPointer()); + return std::make_tuple(activeBatchSize, inputsEmbedsAddr, outputLogitsAddr, loraWeightsName); +} + +trt_edgellm::rt::LLMEngineRunner::BaseGraphKey baseKey(rt::Tensor const& baseTreeDecodingInputsEmbeds, + rt::Tensor const& outputLogits, rt::Tensor const& outputHiddenStates, std::string const& loraWeightsName) noexcept +{ + int64_t const activeBatchSize = baseTreeDecodingInputsEmbeds.getShape()[0]; + uintptr_t const inputsEmbedsAddr = reinterpret_cast(baseTreeDecodingInputsEmbeds.rawPointer()); + uintptr_t const outputLogitsAddr = reinterpret_cast(outputLogits.rawPointer()); + uintptr_t const outputHiddenStatesAddr = reinterpret_cast(outputHiddenStates.rawPointer()); + return std::make_tuple( + activeBatchSize, inputsEmbedsAddr, outputLogitsAddr, outputHiddenStatesAddr, loraWeightsName); +} + +} // namespace + +namespace trt_edgellm +{ +namespace rt +{ + +//! Current implementation limits to two optimization profiles per LLM engine. +static constexpr int32_t kPREFILL_PROFILE_INDEX{0}; +static constexpr int32_t kGENERATION_PROFILE_INDEX{1}; + +bool engineHasIOTensor(nvinfer1::ICudaEngine const* engine, std::string const& name) +{ + for (int32_t i = 0; i < engine->getNbIOTensors(); ++i) + { + char const* tensorName = engine->getIOTensorName(i); + if (tensorName != nullptr && name == tensorName) + { + return true; + } + } + return false; +} + +//! Read output vocab size from logits binding; dual-profile exports may report invalid profile dims. +static int32_t getOutputVocabSizeFromEngine(nvinfer1::ICudaEngine const* engine) +{ + if (!engineHasIOTensor(engine, binding_names::kLogits)) + { + return 0; + } + + auto extractVocabSize = [](nvinfer1::Dims const& dims) -> int32_t { + if (dims.nbDims == 3 && dims.d[2] > 0) + { + return dims.d[2]; + } + if (dims.nbDims == 2 && dims.d[1] > 0) + { + return dims.d[1]; + } + return 0; + }; + + for (int32_t profileIndex : {kPREFILL_PROFILE_INDEX, kGENERATION_PROFILE_INDEX}) + { + int32_t const vocabSize + = extractVocabSize(engine->getProfileShape(binding_names::kLogits, profileIndex, OptProfileSelector::kMAX)); + if (vocabSize > 0) + { + return vocabSize; + } + } + return extractVocabSize(engine->getTensorShape(binding_names::kLogits)); +} + +LLMEngineRunner::LLMEngineRunner(std::filesystem::path const& enginePath, std::filesystem::path const& configPath, + std::unordered_map const& loraWeightsMap, cudaStream_t stream) +{ + + LOG_INFO("Loading config file %s", configPath.string().c_str()); + + // Parse configuration from JSON file + Json configJson; + std::ifstream configFileStream(configPath); + if (!configFileStream.is_open()) + { + LOG_ERROR("Failed to open config file: %s", configPath.string().c_str()); + throw std::runtime_error("Failed to open config file: " + configPath.string()); + } + try + { + configJson = Json::parse(configFileStream); + configFileStream.close(); + } + catch (Json::parse_error const& e) + { + LOG_ERROR("Failed to parse config file with error: %s", e.what()); + throw std::runtime_error("Failed to parse config file: " + configPath.string()); + } + + if (!this->initializeConfigFromJson(configJson)) + { + LOG_ERROR("Failed to initialize LLMEngineRunner from config file: %s", configPath.string().c_str()); + throw std::runtime_error("Failed to initialize LLMEngineRunner from config file: " + configPath.string()); + } + + // Load the engine after config loading succeeds + LOG_INFO("Loading engine file: %s", enginePath.string().c_str()); + mRuntime = std::unique_ptr(nvinfer1::createInferRuntime(gLogger)); + + auto mmapReader = std::make_unique(enginePath); + if (mmapReader->getData() == nullptr) + { + LOG_ERROR("Failed to use MMap to read engine from file path: %s", enginePath.string().c_str()); + throw std::runtime_error("Failed to use MMap to read engine from file path: " + enginePath.string()); + } + mEngine = std::unique_ptr( + mRuntime->deserializeCudaEngine(mmapReader->getData(), mmapReader->getSize())); + + // Use single executionContext for both prefill and generation. + // Context memory is user-managed to enable sharing with other engines. + // The caller must provide shared context memory via setContextMemory() before execution. + mTRTExecutionContext = std::unique_ptr( + mEngine->createExecutionContext(ExecutionContextAllocationStrategy::kUSER_MANAGED)); + + if (trt_edgellm::layerProfiler::LayerProfiler::getInstance().isEnabled()) + { + mTRTExecutionContext->setProfiler(&trt_edgellm::layerProfiler::LayerProfiler::getInstance()); + } + + if (!this->validateConfigFromEngine()) + { + LOG_ERROR("Failed to match config file %s with engine file: %s", configPath.string().c_str(), + enginePath.string().c_str()); + throw std::runtime_error( + "Failed to match config file " + configPath.string() + " with engine file: " + enginePath.string()); + } + + RopeConfig const& ropeConfig = mConfig.ropeConfig; + switch (ropeConfig.type) + { + case RopeType::kLongRope: + { + LOG_DEBUG("Initialize long Rope CosSinCache."); + check::check(ropeConfig.longRope.has_value() && ropeConfig.longRope.value().originalMaxPositionEmbeddings != -1, + "longRope is not set correctly"); + + rt::Tensor shortCosSinCache = rt::Tensor({1, mConfig.maxKVCacheCapacity, mConfig.rotaryDim}, + rt::DeviceType::kGPU, DataType::kFLOAT, "LLMEngineRunner::shortCosSinCache"); + rt::Tensor longCosSinCache = rt::Tensor({1, mConfig.maxKVCacheCapacity, mConfig.rotaryDim}, + rt::DeviceType::kGPU, DataType::kFLOAT, "LLMEngineRunner::longCosSinCache"); + bool const initRopeStatus + = initializeLongRopeCosSinCache(shortCosSinCache, longCosSinCache, ropeConfig, stream); + if (!initRopeStatus) + { + LOG_ERROR("Failed to initialize long Rope CosSinCache."); + throw std::runtime_error("Failed to initialize long Rope CosSinCache."); + } + if (mConfig.maxKVCacheCapacity <= ropeConfig.longRope.value().originalMaxPositionEmbeddings) + { + mPosEncCosSinCache = std::move(shortCosSinCache); + } + else + { + mPosEncCosSinCache = std::move(longCosSinCache); + } + break; + } + case RopeType::kMRope: + { + this->mPosEncCosSinCache + = rt::Tensor({mConfig.maxSupportedBatchSize, mConfig.maxKVCacheCapacity, mConfig.rotaryDim}, + rt::DeviceType::kGPU, DataType::kFLOAT, "LLMEngineRunner::mPosEncCosSinCache"); + + // Initialize MRoPE cache for all batch slots using text-only sequential positions. + kernel::initializeTextOnlyMRopeCosSin(mPosEncCosSinCache.dataPointer(), ropeConfig.rotaryTheta, + mConfig.rotaryDim, mConfig.maxKVCacheCapacity, mConfig.maxSupportedBatchSize, stream); + break; + } + case RopeType::kNoRope: + { + LOG_DEBUG("No RoPE: initializing identity CosSinCache (cos=1, sin=0)."); + this->mPosEncCosSinCache = rt::Tensor({1, mConfig.maxKVCacheCapacity, mConfig.rotaryDim}, rt::DeviceType::kGPU, + DataType::kFLOAT, "LLMEngineRunner::mPosEncCosSinCache"); + bool const initStatus = initializeNopeCosSinCache(mPosEncCosSinCache, stream); + if (!initStatus) + { + LOG_ERROR("Failed to initialize identity CosSinCache."); + throw std::runtime_error("Failed to initialize identity CosSinCache."); + } + break; + } + default: + { + LOG_DEBUG("Initialize persistent Rope CosSinCache."); + this->mPosEncCosSinCache = rt::Tensor({1, mConfig.maxKVCacheCapacity, mConfig.rotaryDim}, rt::DeviceType::kGPU, + DataType::kFLOAT, "LLMEngineRunner::mPosEncCosSinCache"); + bool const initRopeStatus = initializeRopeCosSinCache(mPosEncCosSinCache, ropeConfig, stream); + if (!initRopeStatus) + { + LOG_ERROR("Failed to initialize persistent Rope CosSinCache."); + throw std::runtime_error("Failed to initialize persistent Rope CosSinCache."); + } + break; + } + } + // Bind RopeCosSin cache + bool setRopeCosSinCacheStatus{true}; + setRopeCosSinCacheStatus + &= mTRTExecutionContext->setTensorAddress(binding_names::kRopeCosSin, mPosEncCosSinCache.rawPointer()); + if (!setRopeCosSinCacheStatus) + { + LOG_ERROR("Failed to set rope cos sin cache to the engine"); + throw std::runtime_error("Failed to set rope cos sin cache to the engine"); + } + + if (!validateKVCacheType()) + { + LOG_ERROR("Failed to validate KV cache type"); + throw std::runtime_error("Failed to validate KV cache type"); + } + + // Detect KV cache storage dtype from engine bindings. + nvinfer1::DataType kvCacheType = getKVCacheType(); + + int32_t const kvCacheLayers + = (mConfig.numAttentionLayers > 0) ? mConfig.numAttentionLayers : mConfig.numDecoderLayers; + DataType const recurrentStateType = (mConfig.numLinearAttnLayers > 0) ? getRecurrentStateType() : DataType::kHALF; + DataType const convStateType = (mConfig.numLinearAttnLayers > 0) ? getConvStateType() : DataType::kHALF; + this->mKVCache = rt::LinearKVCache( + rt::LinearKVCache::CacheConfig{kvCacheLayers, mConfig.maxSupportedBatchSize, mConfig.maxKVCacheCapacity, + mConfig.numKVHeads, mConfig.headDim, kvCacheType, + // Recurrent state config (zero for pure-attention models) + mConfig.numLinearAttnLayers, mConfig.recurrentStateNumHeads, mConfig.recurrentStateHeadDim, + mConfig.recurrentStateSize, recurrentStateType, mConfig.convDim, mConfig.convKernel, convStateType}, + stream); + + // Instantiate other GPU memory input that needed by the Engine execution. + this->mSequenceContextLengths = rt::Tensor({mConfig.maxSupportedBatchSize}, rt::DeviceType::kGPU, DataType::kINT32, + "LLMEngineRunner::mSequenceContextLengths"); + CUDA_CHECK( + cudaMemsetAsync(mSequenceContextLengths.rawPointer(), 0, mSequenceContextLengths.getMemoryCapacity(), stream)); + + if (mConfig.enableEagleSpecDecode) + { + // For EAGLE: last_token_ids is 2D [batch_size, num_selected_tokens] to support multi-batch + this->mSelectTokenIndices = rt::Tensor({mConfig.maxSupportedBatchSize, mConfig.maxVerifyTreeSize}, + rt::DeviceType::kGPU, DataType::kINT64, "LLMEngineRunner::mSelectTokenIndices"); + CUDA_CHECK( + cudaMemsetAsync(mSelectTokenIndices.rawPointer(), 0, mSelectTokenIndices.getMemoryCapacity(), stream)); + this->mHostSelectTokenIndices = rt::Tensor({mConfig.maxSupportedBatchSize, mConfig.maxVerifyTreeSize}, + rt::DeviceType::kCPU, DataType::kINT64, "LLMEngineRunner::mHostSelectTokenIndices"); + // Allocate position IDs to support both prefill and tree decoding + int32_t const maxSeqLen = std::max(mConfig.maxSupportedInputLength, mConfig.maxVerifyTreeSize); + this->mEagleBasePositionIds = rt::Tensor({mConfig.maxSupportedBatchSize, maxSeqLen}, rt::DeviceType::kGPU, + DataType::kINT32, "LLMEngineRunner::mEagleBasePositionIds"); + CUDA_CHECK( + cudaMemsetAsync(mEagleBasePositionIds.rawPointer(), 0, mEagleBasePositionIds.getMemoryCapacity(), stream)); + int32_t const packedMaskSize = divUp(mConfig.maxVerifyTreeSize, 32); + this->mEagleBasePackedMask + = rt::Tensor({mConfig.maxSupportedBatchSize, mConfig.maxVerifyTreeSize, packedMaskSize}, + rt::DeviceType::kGPU, DataType::kINT32, "LLMEngineRunner::mEagleBasePackedMask"); + CUDA_CHECK( + cudaMemsetAsync(mEagleBasePackedMask.rawPointer(), 0, mEagleBasePackedMask.getMemoryCapacity(), stream)); + } + else + { + this->mSelectTokenIndices = rt::Tensor({mConfig.maxSupportedBatchSize, 1}, rt::DeviceType::kGPU, + DataType::kINT64, "LLMEngineRunner::mSelectTokenIndices"); + CUDA_CHECK( + cudaMemsetAsync(mSelectTokenIndices.rawPointer(), 0, mSelectTokenIndices.getMemoryCapacity(), stream)); + this->mHostSelectTokenIndices = rt::Tensor({mConfig.maxSupportedBatchSize, 1}, rt::DeviceType::kCPU, + DataType::kINT64, "LLMEngineRunner::mHostSelectTokenIndices"); + } + + // Add the LoRA weights to the engine. + if (isLoraWeightsSupported()) + { + for (auto const& [loraWeightsName, loraWeightsPath] : loraWeightsMap) + { + if (loraWeightsPath.empty()) + { + continue; + } + if (!this->addLoraWeights(loraWeightsName, loraWeightsPath, stream)) + { + LOG_ERROR("Failed to add LoRA weights: %s", loraWeightsName.c_str()); + throw std::runtime_error("Failed to add LoRA weights: " + loraWeightsName); + } + } + } + + // Initialize the dummy tensor as TensorRT does not support nullptr for binding + // Calculate maximum memory requirements across all use cases: + // 1. Attention mask: {maxSupportedBatchSize, 1, 1} + // 2. Attention position IDs: {maxSupportedBatchSize, 1} + // 3. LoRA weights: max dimension across all adapters + // 4. KV cache start index: {maxSupportedBatchSize} + // 5. Deepstack embeds for generation: {maxSupportedBatchSize, maxVerifyTreeSize or 1, hiddenSize} + std::vector dummyInputSizes = { + static_cast( + mConfig.maxSupportedBatchSize), // attention mask/attention position IDs/KV cache start index + static_cast(getMaxLoraWeightsDimension() * kEMPTY_LORA_RANK), // LoRA weights + }; + + // Add deepstack_embeds size for generation profile + // Use maxVerifyTreeSize for eagle or 1 for vanilla decoding + if (mConfig.numDeepstackFeatures > 0) + { + int64_t const deepstackSeqLen = mConfig.enableEagleSpecDecode ? mConfig.maxVerifyTreeSize : 1; + int64_t const deepstackSize + = static_cast(mConfig.maxSupportedBatchSize) * deepstackSeqLen * mConfig.hiddenSize; + dummyInputSizes.push_back(deepstackSize); + } + + int64_t maxDummyElements = *std::max_element(dummyInputSizes.begin(), dummyInputSizes.end()); + mDummyInputTensor = rt::Tensor( + {maxDummyElements}, rt::DeviceType::kGPU, nvinfer1::DataType::kHALF, "LLMEngineRunner::mDummyInputTensor"); + // Initialize dummy tensor memory to zero + CUDA_CHECK(cudaMemsetAsync(mDummyInputTensor.rawPointer(), 0, mDummyInputTensor.getMemoryCapacity(), stream)); + + // Allocate dummy output tensor for hidden_states for Eagle speculative decoding. + // TRT engine under this mode will produce output hidden states. we reserve this buffer to hold the data when + // conduct vanilla decoding. This will make runtime design cleaner. + if (mConfig.enableEagleSpecDecode) + { + int64_t const dummyOutputSize = static_cast(mConfig.maxSupportedBatchSize) * mConfig.outputHiddenDim; + mDummyOutputTensor = rt::Tensor( + {dummyOutputSize}, rt::DeviceType::kGPU, nvinfer1::DataType::kHALF, "LLMEngineRunner::mDummyOutputTensor"); + } + else if (mConfig.enablePrefixKVOutputs) + { + int64_t dummyOutputSize{1}; + for (int64_t const dim : mConfig.prefixKVOutputShape) + { + check::check(dim > 0, "Invalid prefix_k output dimension in config"); + dummyOutputSize *= dim; + } + mDummyOutputTensor = rt::Tensor( + {dummyOutputSize}, rt::DeviceType::kGPU, nvinfer1::DataType::kHALF, "LLMEngineRunner::mDummyOutputTensor"); + } + + if (mConfig.enableContextEmb || mConfig.enableLmHiddenStates) + { + int32_t const auxDim = mConfig.enableContextEmb ? mConfig.contextEmbDim : mConfig.hiddenSize; + int64_t const dummyContextEmbSize + = static_cast(mConfig.maxSupportedBatchSize) * mConfig.maxSupportedInputLength * auxDim; + mDummyContextEmbTensor = rt::Tensor({dummyContextEmbSize}, rt::DeviceType::kGPU, nvinfer1::DataType::kHALF, + "LLMEngineRunner::mDummyContextEmbTensor"); + } + + // Initialize kKVCacheStartIndex to dummy tensor for both profiles to avoid "address not set" error + // when switching optimization profiles. The actual address will be set during runtime execution. + { + bool setKVCacheStartIndexStatus{true}; + Dims const kvStartIndexEngineDim = mEngine->getTensorShape(binding_names::kKVCacheStartIndex); + setKVCacheStartIndexStatus &= mTRTExecutionContext->setTensorAddress( + binding_names::kKVCacheStartIndex, mDummyInputTensor.rawPointer()); + setKVCacheStartIndexStatus + &= mTRTExecutionContext->setInputShape(binding_names::kKVCacheStartIndex, kvStartIndexEngineDim); + if (!setKVCacheStartIndexStatus) + { + LOG_ERROR("Failed to set kKVCacheStartIndex dummy tensor for initialization"); + throw std::runtime_error("Failed to set kKVCacheStartIndex dummy tensor for initialization"); + } + } + // Reset the LoRA weights to zero tensors. + if (!this->resetLoraWeights()) + { + LOG_ERROR("Failed to initialize LoRA weights to zero tensors"); + throw std::runtime_error("Failed to initialize LoRA weights to zero tensors"); + } + + // Synchronize the stream to ensure all the operations have completed. + CUDA_CHECK(cudaStreamSynchronize(stream)); +} + +int64_t LLMEngineRunner::getRequiredContextMemorySize() const +{ + return mEngine->getDeviceMemorySizeV2(); +} + +bool LLMEngineRunner::setContextMemory(rt::Tensor& sharedContextMemory) +{ + int64_t const requiredSize = getRequiredContextMemorySize(); + if (sharedContextMemory.getMemoryCapacity() < requiredSize) + { + LOG_ERROR("Shared context memory (%zu bytes) is smaller than required (%zu bytes)", + static_cast(sharedContextMemory.getMemoryCapacity()), static_cast(requiredSize)); + return false; + } + mTRTExecutionContext->setDeviceMemoryV2(sharedContextMemory.rawPointer(), sharedContextMemory.getMemoryCapacity()); + return true; +} + +nvinfer1::DataType LLMEngineRunner::getKVCacheType() const +{ + if (mConfig.useTrtNativeOps) + { + std::string const trtNativeKVBindingName0 = binding_names::formatKCacheName(/*layerIdx=*/0, /*isPast=*/true); + return mEngine->getTensorDataType(trtNativeKVBindingName0.c_str()); + } + else + { + std::string const pluginKVBindingName0 = binding_names::formatKVCacheName(/*layerIdx=*/0, /*isPast=*/true); + return mEngine->getTensorDataType(pluginKVBindingName0.c_str()); + } +} + +nvinfer1::DataType LLMEngineRunner::getRecurrentStateType() const +{ + std::string const name = binding_names::formatRecurrentStateName(/*recurrentLayerIdx=*/0, /*isPast=*/true); + return mEngine->getTensorDataType(name.c_str()); +} + +nvinfer1::DataType LLMEngineRunner::getConvStateType() const +{ + std::string const name = binding_names::formatConvStateName(/*recurrentLayerIdx=*/0, /*isPast=*/true); + return mEngine->getTensorDataType(name.c_str()); +} + +bool LLMEngineRunner::validateKVCacheType() const +{ + // Sanity check: ensure KV-cache precision (dtype) is consistent across all layers (and both past/present). + // We rely on a single dtype when allocating/owning the KV cache buffers. + if (mConfig.useTrtNativeOps) + { + auto kBindingName0 = binding_names::formatKCacheName(/*layerIdx=*/0, /*isPast=*/true); + DataType const kCacheType0 = mEngine->getTensorDataType(kBindingName0.c_str()); + auto vBindingName0 = binding_names::formatVCacheName(/*layerIdx=*/0, /*isPast=*/true); + auto const checkKVCacheDType = [&](int32_t layerIdx, bool isPast) { + std::string const kBindingName = binding_names::formatKCacheName(layerIdx, isPast); + DataType const kCacheType = mEngine->getTensorDataType(kBindingName.c_str()); + std::string const vBindingName = binding_names::formatVCacheName(layerIdx, isPast); + DataType const vCacheType = mEngine->getTensorDataType(vBindingName.c_str()); + if (kCacheType != kCacheType0 || vCacheType != kCacheType0) + { + LOG_ERROR( + "KV cache dtype mismatch detected. Expected all layers to use the same dtype as '%s' (dtype=%d), " + "but " + "binding '%s' has dtype=%d and '%s' has dtype=%d.", + kBindingName0.c_str(), static_cast(kCacheType0), kBindingName.c_str(), + static_cast(kCacheType), vBindingName.c_str(), static_cast(vCacheType)); + throw std::runtime_error("KV cache dtype mismatch across layers"); + } + }; + int32_t const kvLayers + = (mConfig.numAttentionLayers > 0) ? mConfig.numAttentionLayers : mConfig.numDecoderLayers; + for (int32_t layerIdx = 0; layerIdx < kvLayers; ++layerIdx) + { + checkKVCacheDType(layerIdx, /*isPast=*/true); + std::string const presentKBindingName = binding_names::formatKCacheName(layerIdx, /*isPast=*/false); + if (engineHasIOTensor(mEngine.get(), presentKBindingName)) + { + checkKVCacheDType(layerIdx, /*isPast=*/false); + } + } + } + else + { + auto kvBindingName0 = binding_names::formatKVCacheName(/*layerIdx=*/0, /*isPast=*/true); + DataType const kvCacheType = mEngine->getTensorDataType(kvBindingName0.c_str()); + auto const checkKVCacheDType = [&](int32_t layerIdx, bool isPast) { + std::string const kvBindingName = binding_names::formatKVCacheName(layerIdx, isPast); + DataType const dt = mEngine->getTensorDataType(kvBindingName.c_str()); + if (dt != kvCacheType) + { + LOG_ERROR( + "KV cache dtype mismatch detected. Expected all layers to use the same dtype as '%s' (dtype=%d), " + "but " + "binding '%s' has dtype=%d.", + kvBindingName0.c_str(), static_cast(kvCacheType), kvBindingName.c_str(), + static_cast(dt)); + throw std::runtime_error("KV cache dtype mismatch across layers"); + } + }; + int32_t const kvLayers + = (mConfig.numAttentionLayers > 0) ? mConfig.numAttentionLayers : mConfig.numDecoderLayers; + for (int32_t layerIdx = 0; layerIdx < kvLayers; ++layerIdx) + { + checkKVCacheDType(layerIdx, /*isPast=*/true); + std::string const presentBindingName = binding_names::formatKVCacheName(layerIdx, /*isPast=*/false); + if (engineHasIOTensor(mEngine.get(), presentBindingName)) + { + checkKVCacheDType(layerIdx, /*isPast=*/false); + } + } + } + + return true; +} + +bool LLMEngineRunner::initializeConfigFromJson(Json const& configJson) noexcept +{ + try + { + // Check model version + std::string modelVersion = configJson.value(binding_names::kEdgellmVersion, ""); + version::checkVersion(modelVersion); + + // Define required fields for main config + std::vector const requiredConfigFields + = {"num_hidden_layers", "num_key_value_heads", "head_dim", "vocab_size", "builder_config"}; + + // Validate required fields exist in main config + for (auto const& field : requiredConfigFields) + { + if (!configJson.contains(field)) + { + LOG_ERROR("initializeConfigFromJson(): Missing required field '%s' in config", field.c_str()); + return false; + } + } + + auto const& builderConfig = configJson["builder_config"]; + + // Define required fields for builder_config + std::vector const requiredBuilderConfigFields + = {"max_batch_size", "max_input_len", "max_kv_cache_capacity", "max_lora_rank", "eagle_base"}; + + // Validate required fields exist in builder_config + for (auto const& field : requiredBuilderConfigFields) + { + if (!builderConfig.contains(field)) + { + LOG_ERROR("initializeConfigFromJson(): Missing required field '%s' in builder_config", field.c_str()); + return false; + } + } + + // Extract values with proper type checking + mConfig.numDecoderLayers = configJson["num_hidden_layers"].get(); + mConfig.numKVHeads = configJson["num_key_value_heads"].get(); + mConfig.headDim = configJson["head_dim"].get(); + mConfig.rotaryDim = static_cast(mConfig.headDim * configJson.value("partial_rotary_factor", 1.0f)); + mConfig.hiddenSize = configJson["hidden_size"].get(); + mConfig.vocabSize = configJson["vocab_size"].get(); + // Optional: reduced vocabulary size (0 if not present) + mConfig.reducedVocabSize = configJson.value(binding_names::kReducedVocabSizeKey, 0); + // Set actual output vocab size: use reduced size if enabled, otherwise full size + mConfig.outputVocabSize = (mConfig.reducedVocabSize > 0) ? mConfig.reducedVocabSize : mConfig.vocabSize; + // Read num_deepstack_features if present (Qwen3-VL and Qwen3-Omni models) + mConfig.numDeepstackFeatures = configJson.value("num_deepstack_features", 0); + + // Read audio and image token IDs for Qwen3-Omni (used by embeddingLookupQwen3Omni kernel) + mConfig.audioTokenId = configJson.value("audio_token_id", 0); + mConfig.imageTokenId = configJson.value("image_token_id", 0); + + // Hybrid linear attention configuration (Mamba, GDN, or other linear attention) + mConfig.numLinearAttnLayers = configJson.value("num_linear_attn_layers", 0); + mConfig.numAttentionLayers = configJson.value("num_attention_layers", mConfig.numDecoderLayers); + mConfig.recurrentStateNumHeads = configJson.value("recurrent_state_num_heads", 0); + mConfig.recurrentStateHeadDim = configJson.value("recurrent_state_head_dim", 0); + mConfig.recurrentStateSize = configJson.value("recurrent_state_size", 0); + mConfig.convDim = configJson.value("conv_dim", 0); + mConfig.convKernel = configJson.value("conv_kernel", 0); + + // Extract builder_config values + mConfig.maxSupportedBatchSize = builderConfig["max_batch_size"].get(); + mConfig.maxSupportedInputLength = builderConfig["max_input_len"].get(); + mConfig.maxKVCacheCapacity = builderConfig["max_kv_cache_capacity"].get(); + mConfig.maxSupportedLoraRank = builderConfig["max_lora_rank"].get(); + mConfig.enableEagleSpecDecode = builderConfig["eagle_base"].get(); + mConfig.enableContextEmb = builderConfig.value("context_emb", false); + + // Collect RoPE configuration + mConfig.ropeConfig = collectRopeConfig(configJson); + + // Initialize useTrtNativeOps from builder_config + if (builderConfig.contains("trt_native_ops")) + { + mConfig.useTrtNativeOps = builderConfig["trt_native_ops"].get(); + } + + // Validate configuration values - all must be positive except max_lora_rank + std::vector> positiveFields = {{"num_decoder_layers", mConfig.numDecoderLayers}, + {"num_key_value_heads", mConfig.numKVHeads}, {"head_dim", mConfig.headDim}, + {"rotary_dim", mConfig.rotaryDim}, {"hidden_size", mConfig.hiddenSize}, {"vocab_size", mConfig.vocabSize}, + {"max_batch_size", mConfig.maxSupportedBatchSize}, {"max_input_len", mConfig.maxSupportedInputLength}, + {"max_kv_cache_capacity", mConfig.maxKVCacheCapacity}}; + + for (auto const& [fieldName, value] : positiveFields) + { + if (value <= 0) + { + LOG_ERROR("initializeConfigFromJson(): Invalid %s: %d (must be positive)", fieldName.c_str(), value); + return false; + } + } + + // FIXME: Not a proper way to determine the output hidden dim. + // Hardcode output hidden_dim to 3 x model hidden_size which is default in eagle3. + if (mConfig.enableEagleSpecDecode) + { + mConfig.outputHiddenDim = configJson["hidden_size"].get() * 3; + + // maxVerifyTreeSize is only required when eagle_base is true + if (!builderConfig.contains("max_verify_tree_size")) + { + LOG_ERROR( + "initializeConfigFromJson(): Missing required field 'max_verify_tree_size' in builder_config for " + "Eagle base model"); + return false; + } + mConfig.maxVerifyTreeSize = builderConfig["max_verify_tree_size"].get(); + + // Validate maxVerifyTreeSize (must be positive) + if (mConfig.maxVerifyTreeSize <= 0) + { + LOG_ERROR("initializeConfigFromJson(): Invalid max_verify_tree_size: %d (must be positive)", + mConfig.maxVerifyTreeSize); + return false; + } + } + + if (mConfig.enableContextEmb) + { + if (!configJson.contains("context_hidden_size")) + { + LOG_ERROR( + "initializeConfigFromJson(): Missing required field 'context_hidden_size' in config for " + "context_emb model"); + return false; + } + mConfig.contextEmbDim = configJson["context_hidden_size"].get(); + if (mConfig.contextEmbDim <= 0) + { + LOG_ERROR("initializeConfigFromJson(): Invalid context_hidden_size: %d (must be positive)", + mConfig.contextEmbDim); + return false; + } + } + + // Validate max_lora_rank separately (must be non-negative) + if (mConfig.maxSupportedLoraRank < 0) + { + LOG_ERROR("initializeConfigFromJson(): Invalid max_lora_rank: %d (must be non-negative)", + mConfig.maxSupportedLoraRank); + return false; + } + if (mConfig.maxSupportedInputLength > mConfig.maxKVCacheCapacity) + { + LOG_ERROR( + "initializeConfigFromJson(): Invalid configuration: max_input_len (%d) cannot be greater than " + "max_kv_cache_capacity (%d)", + mConfig.maxSupportedInputLength, mConfig.maxKVCacheCapacity); + return false; + } + + if (configJson.contains("output_names") && configJson.contains("outputs")) + { + auto const& outputNamesJson = configJson.at("output_names"); + auto const& outputsJson = configJson.at("outputs"); + if (outputNamesJson.is_array() && outputsJson.is_array()) + { + for (size_t outputIdx = 0; outputIdx < outputNamesJson.size() && outputIdx < outputsJson.size(); + ++outputIdx) + { + if (outputNamesJson.at(outputIdx).get() != binding_names::kOutputPrefixK) + { + continue; + } + mConfig.enablePrefixKVOutputs = true; + mConfig.prefixKVOutputShape.clear(); + for (auto const& dimJson : outputsJson.at(outputIdx).at("shape")) + { + mConfig.prefixKVOutputShape.push_back(dimJson.get()); + } + break; + } + } + } + } + catch (std::exception const& e) + { + LOG_ERROR("initializeConfigFromJson(): Unexpected error while parsing config: %s", e.what()); + return false; + } + + LOG_INFO("initializeConfigFromJson(): Loaded LLMEngineRunner with config: %s", formatEngineConfig(mConfig).c_str()); + return true; +} + +bool LLMEngineRunner::validateConfigFromEngine() +{ + // Plugin path: combined KV cache [batch, 2, num_kv_heads, seq_len, head_dim] + auto identifyKVCacheBinding = [](std::string const& bindingName, Dims const& tensorDim) { + return tensorDim.nbDims == 5 && bindingName.find(binding_names::kPastKeyValuesTemplate) != std::string::npos; + }; + + // TRT native: separate K cache [batch, num_kv_heads, seq_len, head_dim] + auto identifyTRTNativeKCacheBinding = [](std::string const& bindingName, Dims const& tensorDim) { + return tensorDim.nbDims == 4 && bindingName.find(binding_names::kPresentKCacheTemplate) != std::string::npos; + }; + + // TRT native: separate V cache [batch, num_kv_heads, seq_len, head_dim] + auto identifyTRTNativeVCacheBinding = [](std::string const& bindingName, Dims const& tensorDim) { + return tensorDim.nbDims == 4 && bindingName.find(binding_names::kPresentVCacheTemplate) != std::string::npos; + }; + + // If the engine comes with deepstack embeds binding, it means the engine is Qwen3-VL. + auto identifyDeepstackEmbedsBinding = [](std::string const& bindingName, Dims const& tensorDim) { + return tensorDim.nbDims == 3 && bindingName.find(binding_names::kDeepstackEmbedsTemplate) != std::string::npos; + }; + + auto validate_eq_engine_with_config + = [&](int32_t const& configValue, int32_t const& engineValue, std::string const& name) -> bool { + if (configValue != engineValue) + { + LOG_ERROR("%s is not consistent. From engine: %d, from config: %d", name.c_str(), engineValue, configValue); + return false; + } + return true; + }; + + LOG_DEBUG("Prefill profile info: %s", printEngineInfo(mEngine.get(), kPREFILL_PROFILE_INDEX).c_str()); + LOG_DEBUG("Generation profile info: %s", printEngineInfo(mEngine.get(), kGENERATION_PROFILE_INDEX).c_str()); + + int32_t nbKVCacheInputs{0}; + int32_t nbTRTNativeKCacheInputs{0}; + int32_t nbTRTNativeVCacheInputs{0}; + int32_t nbDeepstackEmbedsInputs{0}; + int32_t numIOBindings = mEngine->getNbIOTensors(); + + // Lambda to validate KV cache dimensions against profile shape + auto validateKVCacheProfile = [&](Dims const& maxKVCacheShape, std::string const& profileName) -> bool { + bool status{true}; + status + &= validate_eq_engine_with_config(mConfig.numKVHeads, maxKVCacheShape.d[2], profileName + ": numKVHeads"); + status &= validate_eq_engine_with_config( + mConfig.maxKVCacheCapacity, maxKVCacheShape.d[3], profileName + ": maxKVCacheCapacity"); + status &= validate_eq_engine_with_config(mConfig.headDim, maxKVCacheShape.d[4], profileName + ": headDim"); + return status; + }; + + bool isOk{true}; + for (int32_t i = 0; i < numIOBindings; ++i) + { + std::string const bindingName = mEngine->getIOTensorName(i); + Dims const tensorDim = mEngine->getTensorShape(bindingName.c_str()); + + if (identifyKVCacheBinding(bindingName, tensorDim)) + { + // Get max profile shapes for both prefill and generation profiles + Dims const maxKVCacheShapePrefill + = mEngine->getProfileShape(bindingName.c_str(), kPREFILL_PROFILE_INDEX, OptProfileSelector::kMAX); + Dims const maxKVCacheShapeGen + = mEngine->getProfileShape(bindingName.c_str(), kGENERATION_PROFILE_INDEX, OptProfileSelector::kMAX); + + // Validate both profiles + isOk &= validateKVCacheProfile(maxKVCacheShapePrefill, "prefill"); + isOk &= validateKVCacheProfile(maxKVCacheShapeGen, "generation"); + ++nbKVCacheInputs; + } + if (identifyDeepstackEmbedsBinding(bindingName, tensorDim)) + { + isOk &= validate_eq_engine_with_config(mConfig.hiddenSize, tensorDim.d[2], "hiddenSize"); + LOG_DEBUG("validateConfigFromEngine(): Found deepstack embeds binding: %s", bindingName.c_str()); + ++nbDeepstackEmbedsInputs; + } + + bool const isTRTNativeKCacheBinding = identifyTRTNativeKCacheBinding(bindingName, tensorDim); + bool const isTRTNativeVCacheBinding = identifyTRTNativeVCacheBinding(bindingName, tensorDim); + if (isTRTNativeKCacheBinding || isTRTNativeVCacheBinding) + { + if (mConfig.numKVHeads != tensorDim.d[1]) + { + LOG_ERROR("numKVHeads is not consistent (TRT native K or V cache). From engine: %d, from config: %d", + tensorDim.d[1], mConfig.numKVHeads); + return false; + } + if (mConfig.maxKVCacheCapacity != tensorDim.d[2]) + { + LOG_ERROR( + "maxSequenceLength is not consistent (TRT native K or V cache). From engine: %d, from config: %d", + tensorDim.d[2], mConfig.maxKVCacheCapacity); + return false; + } + if (mConfig.headDim != tensorDim.d[3]) + { + LOG_ERROR("headDim is not consistent (TRT native K or V cache). From engine: %d, from config: %d", + tensorDim.d[3], mConfig.headDim); + return false; + } + + if (isTRTNativeKCacheBinding) + { + ++nbTRTNativeKCacheInputs; + } + if (isTRTNativeVCacheBinding) + { + ++nbTRTNativeVCacheInputs; + } + } + } + + // Validate KV cache counts based on attention mode + int32_t const expectedKVLayers + = (mConfig.numAttentionLayers > 0) ? mConfig.numAttentionLayers : mConfig.numDecoderLayers; + if (mConfig.useTrtNativeOps) + { + // TRT native mode: expect separate K and V caches + if (nbTRTNativeKCacheInputs != expectedKVLayers) + { + LOG_ERROR("KV cache layer count mismatch (TRT native K cache). From engine: %d, expected: %d", + nbTRTNativeKCacheInputs, expectedKVLayers); + return false; + } + if (nbTRTNativeVCacheInputs != expectedKVLayers) + { + LOG_ERROR("KV cache layer count mismatch (TRT native V cache). From engine: %d, expected: %d", + nbTRTNativeVCacheInputs, expectedKVLayers); + return false; + } + if (nbKVCacheInputs > 0) + { + LOG_ERROR("Found plugin-style KV cache bindings but config specifies TRT native mode"); + return false; + } + } + else + { + // Plugin mode: expect combined KV caches. + if (nbKVCacheInputs != expectedKVLayers) + { + LOG_ERROR( + "KV cache layer count mismatch. From engine: %d, expected: %d", nbKVCacheInputs, expectedKVLayers); + return false; + } + if (nbTRTNativeKCacheInputs > 0 || nbTRTNativeVCacheInputs > 0) + { + LOG_ERROR("Found TRT native-style K/V cache bindings but config specifies plugin attention mode"); + return false; + } + } + isOk &= validate_eq_engine_with_config( + mConfig.numDeepstackFeatures, nbDeepstackEmbedsInputs, "numDeepstackFeatures"); + + Dims const maxInputPrefillShape + = mEngine->getProfileShape(binding_names::kInputsEmbeds, kPREFILL_PROFILE_INDEX, OptProfileSelector::kMAX); + + // inputs_embeds is 3D: [batch_size, seq_len, hidden_size] + isOk &= validate_eq_engine_with_config( + mConfig.maxSupportedInputLength, maxInputPrefillShape.d[1], "maxSupportedInputLength"); + isOk &= validate_eq_engine_with_config(mConfig.hiddenSize, maxInputPrefillShape.d[2], "hiddenSize"); + + // Validate and potentially override maxSupportedBatchSize from engine's actual max profile + int32_t const engineMaxBatchSize = maxInputPrefillShape.d[0]; + isOk &= validate_eq_engine_with_config(mConfig.maxSupportedBatchSize, engineMaxBatchSize, "maxSupportedBatchSize"); + + // Obtain vocab size from logits binding; dual-profile exports may report invalid profile dims. + int32_t const engineOutputVocabSize = getOutputVocabSizeFromEngine(mEngine.get()); + mConfig.hasLogitsOutput = engineHasIOTensor(mEngine.get(), binding_names::kLogits); + if (engineOutputVocabSize <= 0) + { + if (mConfig.hasLogitsOutput) + { + Dims const logitsDim = mEngine->getTensorShape(binding_names::kLogits); + LOG_ERROR("Unexpected logits tensor rank: %d (expected 2 or 3)", logitsDim.nbDims); + isOk = false; + } + else + { + mConfig.outputVocabSize = 1; + } + } + else + { + isOk &= validate_eq_engine_with_config(mConfig.outputVocabSize, engineOutputVocabSize, "outputVocabSize"); + } + + mConfig.hasLastTokenIdsInput = engineHasIOTensor(mEngine.get(), binding_names::kLastTokenIds); + + // Obtain rotary dim from the engine. + Dims const ropeCosSinCacheDim = mEngine->getTensorShape(binding_names::kRopeCosSin); + isOk &= validate_eq_engine_with_config(mConfig.rotaryDim, ropeCosSinCacheDim.d[2], "rotaryDim"); + + if (mConfig.enableContextEmb) + { + Dims const contextEmbedsDim = mEngine->getTensorShape(binding_names::kOutputContextEmbeds); + isOk &= validate_eq_engine_with_config(mConfig.contextEmbDim, contextEmbedsDim.d[2], "contextEmbDim"); + mPrefillAuxOutputName = binding_names::kOutputContextEmbeds; + } + else if (engineHasIOTensor(mEngine.get(), binding_names::kOutputLmHiddenStates)) + { + Dims const lmHiddenDim = mEngine->getTensorShape(binding_names::kOutputLmHiddenStates); + isOk &= validate_eq_engine_with_config(mConfig.hiddenSize, lmHiddenDim.d[2], "lmHiddenDim"); + mConfig.enableLmHiddenStates = true; + mPrefillAuxOutputName = binding_names::kOutputLmHiddenStates; + } + + if (engineHasIOTensor(mEngine.get(), binding_names::kOutputPrefixK) + && engineHasIOTensor(mEngine.get(), binding_names::kOutputPrefixV)) + { + if (!mConfig.enablePrefixKVOutputs || mConfig.prefixKVOutputShape.empty()) + { + LOG_ERROR("Engine exports prefix_k/prefix_v but config.json is missing prefix_k output metadata"); + isOk = false; + } + } + + if (!isOk) + { + LOG_ERROR("Validation failed. Please check the engine configuration."); + } + return isOk; +} + +LLMEngineRunner::~LLMEngineRunner() noexcept +{ + for (auto& [key, graphPair] : mCudaGraphs) + { + cudaGraphDestroy(graphPair.first); + cudaGraphExecDestroy(graphPair.second); + } + for (auto& [key, graphPair] : mBaseTreeDecodingCudaGraphs) + { + cudaGraphDestroy(graphPair.first); + cudaGraphExecDestroy(graphPair.second); + } +} + +bool LLMEngineRunner::bindPluginKVCacheToEngine(int32_t activeBatchSize) +{ + // Prepare special input binding shape for prefill stage KVCache input. + Dims const kvCacheDims = {5, {activeBatchSize, 2, mConfig.numKVHeads, mConfig.maxKVCacheCapacity, mConfig.headDim}}; + int32_t const kvCacheLayers + = (mConfig.numAttentionLayers > 0) ? mConfig.numAttentionLayers : mConfig.numDecoderLayers; + bool status{true}; + // Bind KV cache tensors to execution contexts + for (int32_t i = 0; i < kvCacheLayers; ++i) + { + std::string const pastKeyValuesName = binding_names::formatKVCacheName(i, true); + std::string const presentKeyValuesName = binding_names::formatKVCacheName(i, false); + + rt::Tensor kvCacheBlock = mKVCache.getCombinedKVCacheForDecoderLayer(i); + status &= mTRTExecutionContext->setTensorAddress(pastKeyValuesName.c_str(), kvCacheBlock.rawPointer()); + if (engineHasIOTensor(mEngine.get(), presentKeyValuesName)) + { + status &= mTRTExecutionContext->setTensorAddress(presentKeyValuesName.c_str(), kvCacheBlock.rawPointer()); + } + status &= mTRTExecutionContext->setInputShape(pastKeyValuesName.c_str(), kvCacheDims); + } + return status; +} + +bool LLMEngineRunner::bindTRTNativeKVCacheToEngine(int32_t activeBatchSize) +{ + // TRT native path: separate K and V caches without the "2" dimension + // Shape: [batch, num_kv_heads, seq_len, head_dim] + Dims const kCacheDimIn = {4, {activeBatchSize, mConfig.numKVHeads, mConfig.maxKVCacheCapacity, mConfig.headDim}}; + Dims const vCacheDimIn = {4, {activeBatchSize, mConfig.numKVHeads, mConfig.maxKVCacheCapacity, mConfig.headDim}}; + + int32_t const kvCacheLayers + = (mConfig.numAttentionLayers > 0) ? mConfig.numAttentionLayers : mConfig.numDecoderLayers; + bool status{true}; + // Bind separate K and V cache tensors to execution contexts + for (int32_t i = 0; i < kvCacheLayers; ++i) + { + std::string const pastKCacheName = binding_names::formatKCacheName(i, true); + std::string const presentKCacheName = binding_names::formatKCacheName(i, false); + std::string const pastVCacheName = binding_names::formatVCacheName(i, true); + std::string const presentVCacheName = binding_names::formatVCacheName(i, false); + + std::pair kvCacheBlocks = mKVCache.getSeparateKVCacheForDecoderLayer(i); + rt::Tensor& kCacheBlock = kvCacheBlocks.first; + rt::Tensor& vCacheBlock = kvCacheBlocks.second; + + // Bind K cache + status &= mTRTExecutionContext->setTensorAddress(pastKCacheName.c_str(), kCacheBlock.rawPointer()); + status &= mTRTExecutionContext->setTensorAddress(presentKCacheName.c_str(), kCacheBlock.rawPointer()); + + // Bind V cache + status &= mTRTExecutionContext->setTensorAddress(pastVCacheName.c_str(), vCacheBlock.rawPointer()); + status &= mTRTExecutionContext->setTensorAddress(presentVCacheName.c_str(), vCacheBlock.rawPointer()); + + // Set shapes for K/V cache + status &= mTRTExecutionContext->setInputShape(pastKCacheName.c_str(), kCacheDimIn); + status &= mTRTExecutionContext->setInputShape(pastVCacheName.c_str(), vCacheDimIn); + } + return status; +} + +bool LLMEngineRunner::bindKVCacheToEngine(int32_t activeBatchSize) +{ + if (mConfig.useTrtNativeOps) + { + return bindTRTNativeKVCacheToEngine(activeBatchSize); + } + else + { + return bindPluginKVCacheToEngine(activeBatchSize); + } +} + +bool LLMEngineRunner::bindRecurrentStateToEngine(int32_t activeBatchSize) +{ + if (mConfig.numLinearAttnLayers == 0) + { + return true; + } + + Dims const recurrentStateDims = {4, + {activeBatchSize, mConfig.recurrentStateNumHeads, mConfig.recurrentStateHeadDim, mConfig.recurrentStateSize}}; + bool status{true}; + for (int32_t i = 0; i < mConfig.numLinearAttnLayers; ++i) + { + rt::Tensor recurrentState = mKVCache.getRecurrentStateForLayer(i); + std::string const pastName = binding_names::formatRecurrentStateName(i, /*isPast=*/true); + std::string const presentName = binding_names::formatRecurrentStateName(i, /*isPast=*/false); + + status &= mTRTExecutionContext->setTensorAddress(pastName.c_str(), recurrentState.rawPointer()); + status &= mTRTExecutionContext->setTensorAddress(presentName.c_str(), recurrentState.rawPointer()); + status &= mTRTExecutionContext->setInputShape(pastName.c_str(), recurrentStateDims); + } + return status; +} + +bool LLMEngineRunner::bindConvStateToEngine(int32_t activeBatchSize) +{ + if (mConfig.numLinearAttnLayers == 0) + { + return true; + } + + Dims const convStateDims = {3, {activeBatchSize, mConfig.convDim, mConfig.convKernel}}; + bool status{true}; + for (int32_t i = 0; i < mConfig.numLinearAttnLayers; ++i) + { + rt::Tensor convState = mKVCache.getConvStateForLayer(i); + std::string const pastName = binding_names::formatConvStateName(i, /*isPast=*/true); + std::string const presentName = binding_names::formatConvStateName(i, /*isPast=*/false); + + status &= mTRTExecutionContext->setTensorAddress(pastName.c_str(), convState.rawPointer()); + status &= mTRTExecutionContext->setTensorAddress(presentName.c_str(), convState.rawPointer()); + status &= mTRTExecutionContext->setInputShape(pastName.c_str(), convStateDims); + } + return status; +} + +rt::Tensor& LLMEngineRunner::getRopeCosSinCacheTensor() noexcept +{ + return mPosEncCosSinCache; +} + +LLMEngineRunnerConfig LLMEngineRunner::getEngineConfig() const noexcept +{ + return mConfig; +} + +rt::LinearKVCache& LLMEngineRunner::getLinearKVCache() noexcept +{ + return mKVCache; +} + +bool LLMEngineRunner::setLMHeadWeights(std::string const& name, rt::Tensor const& tensor) +{ + bool status = mTRTExecutionContext->setTensorAddress(name.c_str(), const_cast(tensor.rawPointer())); + if (!status) + { + LOG_ERROR("setTensorAddress failed for '%s'", name.c_str()); + return false; + } + + bool shapeStatus = mTRTExecutionContext->setInputShape(name.c_str(), tensor.getShape().getTRTDims()); + if (!shapeStatus) + { + LOG_ERROR( + "setInputShape failed for '%s' with shape %s", name.c_str(), tensor.getShape().formatString().c_str()); + return false; + } + + return true; +} + +bool LLMEngineRunner::prefillStepInputValidation(rt::Tensor const& inputsEmbeds, rt::Tensor const& contextLengths, + rt::Tensor const& outputLogits, OptionalOutputTensor outputHiddenStates, OptionalOutputTensor outputContextEmbeds, + rt::OptionalInputTensors deepstackEmbeds) noexcept +{ + int32_t activeBatchSize = inputsEmbeds.getShape()[0]; + int32_t prefillSequenceLength = inputsEmbeds.getShape()[1]; + + // Validate inputsEmbeds + bool const checkInputsGPUTensor = inputsEmbeds.getDeviceType() == rt::DeviceType::kGPU + && inputsEmbeds.getDataType() == nvinfer1::DataType::kHALF && inputsEmbeds.getShape().getNumDims() == 3 + && inputsEmbeds.getShape()[2] == mConfig.hiddenSize && contextLengths.getDeviceType() == rt::DeviceType::kCPU + && outputLogits.getDeviceType() == rt::DeviceType::kGPU; + if (!checkInputsGPUTensor) + { + LOG_ERROR( + "Invalid device type or shape of I/O tensors. InputsEmbeds should be 3D FLOAT16 on GPU with shape " + "[batchSize, seqLen, %d], " + "ContextLengths input should reside on CPU and the rest should reside on GPU.", + mConfig.hiddenSize); + return false; + } + bool const isBatchValid = activeBatchSize <= mConfig.maxSupportedBatchSize + && contextLengths.getShape()[0] == activeBatchSize + && (!mConfig.hasLogitsOutput || outputLogits.getShape()[0] == activeBatchSize); + if (!isBatchValid) + { + LOG_ERROR( + "Invalid batchSize of the input tensors. Either batchSize is larger than " + "maxSupportedBatchSize or batchSize is not consistent among the input tensors. " + "Current inputsEmbeds shape: %s, contextLengths shape: %s, logits shape: %s", + inputsEmbeds.getShape().formatString().c_str(), contextLengths.getShape().formatString().c_str(), + outputLogits.getShape().formatString().c_str()); + return false; + } + if (prefillSequenceLength > mConfig.maxSupportedInputLength) + { + LOG_ERROR( + "Invalid sequence length of the input tensors. Input sequence length (%d) is larger " + "than maxSupportedInputLength (%d). Current inputsEmbeds shape: %s.", + prefillSequenceLength, mConfig.maxSupportedInputLength, inputsEmbeds.getShape().formatString().c_str()); + return false; + } + + // Validate deepstack embeds for Qwen3-VL (these are already embedded) + int32_t deepstackEmbedsCount = static_cast(deepstackEmbeds.size()); + if ((deepstackEmbedsCount != mConfig.numDeepstackFeatures) && (deepstackEmbedsCount != 0)) + { + LOG_ERROR("Invalid deepstack embeds count. Expected either %d or 0, got %d", mConfig.numDeepstackFeatures, + deepstackEmbedsCount); + return false; + } + + // Validate each deepstack embed tensor + for (int32_t i = 0; i < deepstackEmbedsCount; ++i) + { + rt::Tensor const& tensor = deepstackEmbeds[i].get(); + bool const isTensorValid = tensor.getDeviceType() == rt::DeviceType::kGPU && tensor.getShape().getNumDims() == 3 + && tensor.getShape()[0] == activeBatchSize && tensor.getShape()[1] == prefillSequenceLength + && tensor.getShape()[2] == mConfig.hiddenSize; + if (!isTensorValid) + { + LOG_ERROR( + "Invalid deepstack embed at index %d. Expected device type: GPU, shape: [%d, %d, %d]. Current shape: " + "%s", + i, activeBatchSize, prefillSequenceLength, mConfig.hiddenSize, + tensor.getShape().formatString().c_str()); + return false; + } + } + + bool const isLogitsShapeValid = !mConfig.hasLogitsOutput + || (outputLogits.getShape().getNumDims() == 2 && outputLogits.getShape()[1] == mConfig.outputVocabSize); + if (!isLogitsShapeValid) + { + LOG_ERROR( + "Invalid shape of the output logits tensor. The output logits tensor should have shape " + "[activeBatchSize, outputVocabSize]. Current logits shape is %s.", + outputLogits.getShape().formatString().c_str()); + return false; + } + if (mConfig.enableEagleSpecDecode) + { + bool const isHiddenStatesShapeValid = outputHiddenStates.has_value() + && outputHiddenStates.value().get().getShape().getNumDims() == 3 + && outputHiddenStates.value().get().getShape()[0] == activeBatchSize + && outputHiddenStates.value().get().getShape()[1] == prefillSequenceLength + && outputHiddenStates.value().get().getShape()[2] == mConfig.outputHiddenDim; + if (!isHiddenStatesShapeValid) + { + LOG_ERROR( + "With SpecDecode enabled, the output hidden states tensor shall be valid and has shape " + "[activeBatchSize, %d, %d]. Current hidden states shape is %s.", + prefillSequenceLength, mConfig.outputHiddenDim, + outputHiddenStates.value().get().getShape().formatString().c_str()); + return false; + } + } + if (mConfig.enableContextEmb || mConfig.enableLmHiddenStates) + { + int32_t const auxDim = mConfig.enableContextEmb ? mConfig.contextEmbDim : mConfig.hiddenSize; + char const* auxName + = mConfig.enableContextEmb ? binding_names::kOutputContextEmbeds : binding_names::kOutputLmHiddenStates; + bool const isContextEmbedsShapeValid = outputContextEmbeds.has_value() + && outputContextEmbeds.value().get().getShape().getNumDims() == 3 + && outputContextEmbeds.value().get().getShape()[0] == activeBatchSize + && outputContextEmbeds.value().get().getShape()[1] == prefillSequenceLength + && outputContextEmbeds.value().get().getShape()[2] == auxDim; + if (!isContextEmbedsShapeValid) + { + LOG_ERROR( + "With %s enabled, the output %s tensor shall be valid and has shape " + "[activeBatchSize, %d, %d]. Current tensor shape is %s.", + mConfig.enableContextEmb ? "context_emb" : "lm_hidden_states", auxName, prefillSequenceLength, auxDim, + outputContextEmbeds.has_value() ? outputContextEmbeds.value().get().getShape().formatString().c_str() + : "nullopt"); + return false; + } + } + + return true; +} + +bool LLMEngineRunner::executePrefillStep(rt::Tensor const& inputsEmbeds, rt::Tensor const& hostContextLengths, + rt::OptionalInputTensors deepstackEmbeds, rt::Tensor& outputLogits, rt::OptionalOutputTensor outputHiddenStates, + cudaStream_t stream, rt::OptionalOutputTensor outputContextEmbeds, rt::OptionalOutputTensor outputPrefixK, + rt::OptionalOutputTensor outputPrefixV) +{ + bool setOptimizationProfileStatus{true}; + setOptimizationProfileStatus &= mTRTExecutionContext->setOptimizationProfileAsync(kPREFILL_PROFILE_INDEX, stream); + if (!setOptimizationProfileStatus) + { + LOG_ERROR("Failed to set optimization profile to the engine"); + throw std::runtime_error("Failed to set optimization profile to the engine"); + } + + bool const validateInputStatus = this->prefillStepInputValidation( + inputsEmbeds, hostContextLengths, outputLogits, outputHiddenStates, outputContextEmbeds, deepstackEmbeds); + if (!validateInputStatus) + { + LOG_ERROR("executePrefill(): Prefill request not performed due to invalid input tensors."); + return false; + } + + // Verify input tensorShape is valid. + int32_t activeBatchSize = inputsEmbeds.getShape()[0]; + + bool reshapeStatus{true}; + // conduct preparation work for the engine execution. Provide correct shapes for MISC input tensors. + // All models (EAGLE and vanilla) now use 2D shape [batch_size, num_tokens] for last_token_ids + reshapeStatus &= mSelectTokenIndices.reshape({activeBatchSize, 1}); + reshapeStatus &= mSequenceContextLengths.reshape({activeBatchSize}); + if (!reshapeStatus) + { + LOG_ERROR("Failed to reshape select token indices and sequence context lengths for prefill step."); + return false; + } + + check::check(mHostSelectTokenIndices.reshape({activeBatchSize, 1}), "Tensor reshape failed"); + int64_t* selectTokenIndicesData = mHostSelectTokenIndices.dataPointer(); + int32_t const* contextLengthsData = hostContextLengths.dataPointer(); + for (int32_t i = 0; i < activeBatchSize; ++i) + { + selectTokenIndicesData[i] = static_cast(contextLengthsData[i] - 1); + } + CUDA_CHECK(cudaMemcpyAsync(mSelectTokenIndices.rawPointer(), mHostSelectTokenIndices.rawPointer(), + activeBatchSize * sizeof(int64_t), cudaMemcpyHostToDevice, stream)); + CUDA_CHECK(cudaMemcpyAsync(mSequenceContextLengths.rawPointer(), hostContextLengths.rawPointer(), + activeBatchSize * sizeof(int32_t), cudaMemcpyHostToDevice, stream)); + + bool setEngineIOStatus{true}; + // Engine input tensors - bind inputs_embeds directly + setEngineIOStatus &= mTRTExecutionContext->setTensorAddress( + binding_names::kInputsEmbeds, const_cast(inputsEmbeds.rawPointer())); + setEngineIOStatus + &= mTRTExecutionContext->setInputShape(binding_names::kInputsEmbeds, inputsEmbeds.getShape().getTRTDims()); + setEngineIOStatus + &= mTRTExecutionContext->setTensorAddress(binding_names::kContextLengths, mSequenceContextLengths.rawPointer()); + setEngineIOStatus &= mTRTExecutionContext->setInputShape( + binding_names::kContextLengths, mSequenceContextLengths.getShape().getTRTDims()); + if (mConfig.hasLastTokenIdsInput) + { + setEngineIOStatus + &= mTRTExecutionContext->setTensorAddress(binding_names::kLastTokenIds, mSelectTokenIndices.rawPointer()); + setEngineIOStatus &= mTRTExecutionContext->setInputShape( + binding_names::kLastTokenIds, mSelectTokenIndices.getShape().getTRTDims()); + } + + // Setup the KVCache start index tensor. If all KVCache are empty then we can supply zero tensor to the engine. + // Otherwise, we shall supply the KVCache lengths tensor to the engine. + if (!mConfig.useTrtNativeOps && mKVCache.getKVCacheAllEmpty()) + { + setEngineIOStatus &= mTRTExecutionContext->setTensorAddress( + binding_names::kKVCacheStartIndex, mDummyInputTensor.rawPointer()); + setEngineIOStatus + &= mTRTExecutionContext->setInputShape(binding_names::kKVCacheStartIndex, rt::Coords{0}.getTRTDims()); + } + else + { + setEngineIOStatus &= mTRTExecutionContext->setTensorAddress( + binding_names::kKVCacheStartIndex, mKVCache.getKVCacheLengths().rawPointer()); + setEngineIOStatus &= mTRTExecutionContext->setInputShape( + binding_names::kKVCacheStartIndex, mKVCache.getKVCacheLengths().getShape().getTRTDims()); + } + + // RopeCosSin tensor address is set during object construction. We only set shape here to accommodate ND-Rope. + // For ND-RoPE like MRope, reshape the RopeCosSinCache to match the activeBatchSize + if (mConfig.ropeConfig.type == RopeType::kMRope) + { + check::check(mPosEncCosSinCache.reshape({activeBatchSize, mConfig.maxKVCacheCapacity, mConfig.rotaryDim}), + "Tensor reshape failed"); + } + + // For persistent rope, the cache is fixed at {1, maxSeqLen, rotaryDim} and shared across all batches. + setEngineIOStatus + &= mTRTExecutionContext->setInputShape(binding_names::kRopeCosSin, mPosEncCosSinCache.getShape().getTRTDims()); + + // Process deepstack embeds: bind already-embedded tensors to engine + // Runtime must provide these tensors (zero tensors for non-multimodal use cases) + if (mConfig.numDeepstackFeatures > 0) + { + // Bind deepstack embeds to engine by index + for (int32_t idx = 0; idx < mConfig.numDeepstackFeatures; ++idx) + { + rt::Tensor const& embedTensor = deepstackEmbeds[idx].get(); + std::string embedName = binding_names::formatDeepstackEmbedsName(idx); + + setEngineIOStatus &= mTRTExecutionContext->setTensorAddress( + embedName.c_str(), const_cast(embedTensor.rawPointer())); + setEngineIOStatus + &= mTRTExecutionContext->setInputShape(embedName.c_str(), embedTensor.getShape().getTRTDims()); + } + } + // Bind hidden states output if requested + if (outputHiddenStates.has_value()) + { + setEngineIOStatus &= mTRTExecutionContext->setTensorAddress( + binding_names::kOutputHiddenStates, outputHiddenStates.value().get().rawPointer()); + } + + if (mConfig.enablePrefixKVOutputs) + { + if (outputPrefixK.has_value()) + { + setEngineIOStatus &= mTRTExecutionContext->setTensorAddress( + binding_names::kOutputPrefixK, outputPrefixK.value().get().rawPointer()); + } + else if (mDummyOutputTensor.getMemoryCapacity() > 0) + { + setEngineIOStatus &= mTRTExecutionContext->setTensorAddress( + binding_names::kOutputPrefixK, mDummyOutputTensor.rawPointer()); + } + if (outputPrefixV.has_value()) + { + setEngineIOStatus &= mTRTExecutionContext->setTensorAddress( + binding_names::kOutputPrefixV, outputPrefixV.value().get().rawPointer()); + } + else if (mDummyOutputTensor.getMemoryCapacity() > 0) + { + setEngineIOStatus &= mTRTExecutionContext->setTensorAddress( + binding_names::kOutputPrefixV, mDummyOutputTensor.rawPointer()); + } + } + + if (!mPrefillAuxOutputName.empty()) + { + if (outputContextEmbeds.has_value()) + { + setEngineIOStatus &= mTRTExecutionContext->setTensorAddress( + mPrefillAuxOutputName.c_str(), outputContextEmbeds.value().get().rawPointer()); + } + else if (mDummyContextEmbTensor.getMemoryCapacity() > 0) + { + setEngineIOStatus &= mTRTExecutionContext->setTensorAddress( + mPrefillAuxOutputName.c_str(), mDummyContextEmbTensor.rawPointer()); + } + } + + if (mConfig.enableEagleSpecDecode) + { + // Mask input and optional token pos-ids are not used, set to dummy data. + setEngineIOStatus + &= mTRTExecutionContext->setTensorAddress(binding_names::kAttentionMask, mDummyInputTensor.rawPointer()); + setEngineIOStatus &= mTRTExecutionContext->setInputShape( + binding_names::kAttentionMask, Coords{activeBatchSize, 1, 1}.getTRTDims()); + setEngineIOStatus + &= mTRTExecutionContext->setTensorAddress(binding_names::kAttentionPosId, mDummyInputTensor.rawPointer()); + setEngineIOStatus &= mTRTExecutionContext->setInputShape( + binding_names::kAttentionPosId, Coords{activeBatchSize, 1}.getTRTDims()); + } + + // Engine output tensors. + if (mConfig.hasLogitsOutput) + { + setEngineIOStatus &= mTRTExecutionContext->setTensorAddress(binding_names::kLogits, outputLogits.rawPointer()); + } + + // Bind the KVCache IO to the engine + setEngineIOStatus &= this->bindKVCacheToEngine(activeBatchSize); + + // Bind recurrent state for hybrid layers + setEngineIOStatus &= this->bindConvStateToEngine(activeBatchSize); + setEngineIOStatus &= this->bindRecurrentStateToEngine(activeBatchSize); + + if (!setEngineIOStatus) + { + LOG_ERROR("executePrefill(): Failed to bind engine input and output tensors."); + return false; + } + + // launch the engine execution. + bool executeStatus{true}; + executeStatus &= mTRTExecutionContext->enqueueV3(stream); + if (!executeStatus) + { + LOG_ERROR("executePrefill(): Failed on TensorRT prefill stage enqueueV3() call."); + return false; + } + // Prefill operation has completed, commit the new contents with KVCache. + mKVCache.commitSequenceLength(mSequenceContextLengths, stream); + + LOG_DEBUG("executePrefill(): Prefill stage execution completed for request with batch size %d.", activeBatchSize); + return true; +} + +bool LLMEngineRunner::vanillaDecodingStepInputValidation( + rt::Tensor const& inputsEmbeds, rt::Tensor const& outputLogits) noexcept +{ + int32_t activeBatchSize = inputsEmbeds.getShape()[0]; + bool const checkInputsGPUTensor = inputsEmbeds.getDeviceType() == rt::DeviceType::kGPU + && inputsEmbeds.getDataType() == nvinfer1::DataType::kHALF + && outputLogits.getDeviceType() == rt::DeviceType::kGPU; + if (!checkInputsGPUTensor) + { + LOG_ERROR( + "Invalid device type of the input tensors. inputsEmbeds (FLOAT16) and outputLogits " + "should reside on GPU."); + return false; + } + int32_t activeKVCacheBatchSize = mKVCache.getActiveBatchSize(); + bool const isBatchValid = activeBatchSize == activeKVCacheBatchSize; + if (!isBatchValid) + { + LOG_ERROR( + "Invalid batchSize of the input tensors. batchSize shall be equal to the active batch " + "size set by the previous prefill stage."); + return false; + } + bool checkInputShapeValid = inputsEmbeds.getShape().getNumDims() == 3 && inputsEmbeds.getShape()[1] == 1 + && inputsEmbeds.getShape()[2] == mConfig.hiddenSize && outputLogits.getShape().getNumDims() == 2 + && outputLogits.getShape()[1] == mConfig.outputVocabSize; + if (!checkInputShapeValid) + { + LOG_ERROR( + "Invalid shape of the input tensors. The input tensor should have shape " + "[activeBatchSize, 1, hiddenSize] and the output tensor should have shape [activeBatchSize, " + "outputVocabSize]."); + return false; + } + + return true; +} + +bool LLMEngineRunner::vanillaDecodingStepPrepareInputs(int32_t activeBatchSize, cudaStream_t stream) +{ + // For vanilla decode stage, the selected token indices are always 0. + // Also setup the sequence length of each sequence for this run based on committed KVCache length. + check::check(mSelectTokenIndices.reshape({activeBatchSize, 1}), "Tensor reshape failed"); + check::check(mSequenceContextLengths.reshape({activeBatchSize}), "Tensor reshape failed"); + + // For MRope (VLM), reshape the RopeCosSinCache to match the activeBatchSize + if (mConfig.ropeConfig.type == RopeType::kMRope) + { + check::check(mPosEncCosSinCache.reshape({activeBatchSize, mConfig.maxKVCacheCapacity, mConfig.rotaryDim}), + "Tensor reshape failed"); + } + + CUDA_CHECK(cudaMemsetAsync(mSelectTokenIndices.rawPointer(), 0, activeBatchSize * sizeof(int64_t), stream)); + + // For TRT native path, the sequence length input always refer to the length of Q. + // For plugin path, the sequence length refer to the length of K and V. + if (mConfig.useTrtNativeOps) + { + CUDA_CHECK(cudaMemsetAsync(mSequenceContextLengths.rawPointer(), 0, activeBatchSize * sizeof(int32_t), stream)); + } + else + { + // Get KV cache lengths + rt::Tensor& kvCacheLengths = mKVCache.getKVCacheLengths(); + CUDA_CHECK(cudaMemcpyAsync(mSequenceContextLengths.rawPointer(), kvCacheLengths.rawPointer(), + activeBatchSize * sizeof(int32_t), cudaMemcpyDeviceToDevice, stream)); + } + + // Increment the sequence length due to the implementation constraint of AttentionPlugin. + constexpr int32_t kDECODE_INCREMENT{1}; + kernel::incrementLengthTensor(mSequenceContextLengths, kDECODE_INCREMENT, stream); + + return true; +} + +bool LLMEngineRunner::vanillaDecodingStepBindTensors(rt::Tensor const& inputsEmbeds, rt::Tensor& outputLogits, + rt::OptionalOutputTensor outputHiddenStates, rt::OptionalOutputTensor outputContextEmbeds, int32_t activeBatchSize) +{ + bool setEngineIOStatus{true}; + // Engine input tensors - bind inputs_embeds directly + setEngineIOStatus &= mTRTExecutionContext->setTensorAddress( + binding_names::kInputsEmbeds, const_cast(inputsEmbeds.rawPointer())); + setEngineIOStatus + &= mTRTExecutionContext->setInputShape(binding_names::kInputsEmbeds, inputsEmbeds.getShape().getTRTDims()); + setEngineIOStatus + &= mTRTExecutionContext->setTensorAddress(binding_names::kContextLengths, mSequenceContextLengths.rawPointer()); + setEngineIOStatus &= mTRTExecutionContext->setInputShape( + binding_names::kContextLengths, mSequenceContextLengths.getShape().getTRTDims()); + setEngineIOStatus + &= mTRTExecutionContext->setTensorAddress(binding_names::kLastTokenIds, mSelectTokenIndices.rawPointer()); + setEngineIOStatus &= mTRTExecutionContext->setInputShape( + binding_names::kLastTokenIds, mSelectTokenIndices.getShape().getTRTDims()); + setEngineIOStatus &= mTRTExecutionContext->setTensorAddress( + binding_names::kKVCacheStartIndex, mKVCache.getKVCacheLengths().rawPointer()); + setEngineIOStatus &= mTRTExecutionContext->setInputShape( + binding_names::kKVCacheStartIndex, mKVCache.getKVCacheLengths().getShape().getTRTDims()); + + setEngineIOStatus + &= mTRTExecutionContext->setInputShape(binding_names::kRopeCosSin, mPosEncCosSinCache.getShape().getTRTDims()); + + // Update KV cache shapes to match activeBatchSize + setEngineIOStatus &= this->bindKVCacheToEngine(activeBatchSize); + setEngineIOStatus &= this->bindConvStateToEngine(activeBatchSize); + setEngineIOStatus &= this->bindRecurrentStateToEngine(activeBatchSize); + + // Bind deepstack_embeds to dummy tensors for Qwen3VL models during decoding + if (mConfig.numDeepstackFeatures > 0) + { + for (int32_t idx = 0; idx < mConfig.numDeepstackFeatures; ++idx) + { + std::string deepstackEmbedName = binding_names::formatDeepstackEmbedsName(idx); + setEngineIOStatus + &= mTRTExecutionContext->setTensorAddress(deepstackEmbedName.c_str(), mDummyInputTensor.rawPointer()); + setEngineIOStatus &= mTRTExecutionContext->setInputShape( + deepstackEmbedName.c_str(), rt::Coords{activeBatchSize, 1, mConfig.hiddenSize}.getTRTDims()); + } + } + + // Engine output tensors. + if (mConfig.hasLogitsOutput) + { + setEngineIOStatus &= mTRTExecutionContext->setTensorAddress(binding_names::kLogits, outputLogits.rawPointer()); + } + + if (outputHiddenStates.has_value()) + { + setEngineIOStatus &= mTRTExecutionContext->setTensorAddress( + binding_names::kOutputHiddenStates, outputHiddenStates.value().get().rawPointer()); + } + else if (mDummyOutputTensor.getMemoryCapacity() > 0 + && engineHasIOTensor(mEngine.get(), binding_names::kOutputHiddenStates)) + { + // Engine has hidden_states output but user doesn't need it + // Bind to dummy buffer (EAGLE or Qwen3-Omni text-only mode) + setEngineIOStatus &= mTRTExecutionContext->setTensorAddress( + binding_names::kOutputHiddenStates, mDummyOutputTensor.rawPointer()); + } + + if (mConfig.enablePrefixKVOutputs && mDummyOutputTensor.getMemoryCapacity() > 0) + { + if (engineHasIOTensor(mEngine.get(), binding_names::kOutputPrefixK)) + { + setEngineIOStatus &= mTRTExecutionContext->setTensorAddress( + binding_names::kOutputPrefixK, mDummyOutputTensor.rawPointer()); + } + if (engineHasIOTensor(mEngine.get(), binding_names::kOutputPrefixV)) + { + setEngineIOStatus &= mTRTExecutionContext->setTensorAddress( + binding_names::kOutputPrefixV, mDummyOutputTensor.rawPointer()); + } + } + + if (!mPrefillAuxOutputName.empty()) + { + if (outputContextEmbeds.has_value()) + { + setEngineIOStatus &= mTRTExecutionContext->setTensorAddress( + mPrefillAuxOutputName.c_str(), outputContextEmbeds.value().get().rawPointer()); + } + else if (mDummyContextEmbTensor.getMemoryCapacity() > 0) + { + setEngineIOStatus &= mTRTExecutionContext->setTensorAddress( + mPrefillAuxOutputName.c_str(), mDummyContextEmbTensor.rawPointer()); + } + } + + return setEngineIOStatus; +} + +bool LLMEngineRunner::executeVanillaDecodingStep(rt::Tensor const& inputsEmbeds, rt::Tensor& outputLogits, + rt::OptionalOutputTensor outputHiddenStates, cudaStream_t stream, rt::OptionalOutputTensor outputContextEmbeds) +{ + bool const validateInputStatus = this->vanillaDecodingStepInputValidation(inputsEmbeds, outputLogits); + if (!validateInputStatus) + { + LOG_ERROR("executeGeneration(): Generation request not performed due to invalid input tensors."); + return false; + } + + int32_t const activeBatchSize = inputsEmbeds.getShape()[0]; + if (!vanillaDecodingStepPrepareInputs(activeBatchSize, stream)) + { + LOG_ERROR("Failed to prepare inputs."); + return false; + } + + // Launch cuda graph if available for this request, otherwise proceed with normal TensorRT engine execution step. + auto const graphHash = decodingKey(inputsEmbeds, outputLogits, mActiveLoraWeightsName); + if (mCudaGraphs.find(graphHash) != mCudaGraphs.end()) + { + LOG_DEBUG("Use pre-captured CUDA graph for vanilla decoding step."); + cudaGraphExec_t graphExec = mCudaGraphs[graphHash].second; + CUDA_CHECK(cudaGraphLaunch(graphExec, stream)); + } + else + { + bool setOptimizationProfileStatus{true}; + setOptimizationProfileStatus + &= mTRTExecutionContext->setOptimizationProfileAsync(kGENERATION_PROFILE_INDEX, stream); + if (!setOptimizationProfileStatus) + { + LOG_ERROR("Failed to set optimization profile to the engine"); + throw std::runtime_error("Failed to set optimization profile to the engine"); + } + + LOG_INFO("Vanilla decoding step CUDA graph not captured."); + if (!vanillaDecodingStepBindTensors( + inputsEmbeds, outputLogits, outputHiddenStates, outputContextEmbeds, activeBatchSize)) + { + LOG_ERROR("Failed to bind tensors."); + return false; + } + + // launch the engine execution. + bool executeStatus{true}; + executeStatus &= mTRTExecutionContext->enqueueV3(stream); + if (!executeStatus) + { + LOG_ERROR("Failed on TensorRT decode stage enqueueV3() call."); + return false; + } + } + + // Completed decoding step, commit the KVCache length of this run. + constexpr int32_t kVANILLA_DECODE_INCREMENT{1}; + mKVCache.commitSequenceLength(kVANILLA_DECODE_INCREMENT, stream); + LOG_DEBUG("Decoding stage execution completed for request with batch size %d.", activeBatchSize); + return true; +} + +bool LLMEngineRunner::eagleBaseTreeDecodingStepInputValidation(rt::Tensor const& baseTreeDecodingInputsEmbeds, + rt::Tensor const& baseTreeDecodingMask, rt::Tensor const& outputLogits, + rt::Tensor const& outputHiddenStates) noexcept +{ + // All input tensors shall reside on GPU. + bool const checkInputsGPUTensor = baseTreeDecodingInputsEmbeds.getDeviceType() == rt::DeviceType::kGPU + && baseTreeDecodingMask.getDeviceType() == rt::DeviceType::kGPU + && outputLogits.getDeviceType() == rt::DeviceType::kGPU + && outputHiddenStates.getDeviceType() == rt::DeviceType::kGPU; + if (!checkInputsGPUTensor) + { + LOG_ERROR( + "eagleBaseTreeDecodingStepInputValidation(): Invalid device type of I/O tensors. All inputs and outputs " + "shall " + "reside on GPU."); + return false; + } + // Validate datatypes of the input tensors. + bool const isInputTypeValid = baseTreeDecodingInputsEmbeds.getDataType() == DataType::kHALF + && baseTreeDecodingMask.getDataType() == DataType::kINT8 && outputLogits.getDataType() == DataType::kFLOAT + && outputHiddenStates.getDataType() == DataType::kHALF; + if (!isInputTypeValid) + { + LOG_ERROR( + "eagleBaseTreeDecodingStepInputValidation(): Input embeds shall be FLOAT16, hidden states I/O shall be " + "FLOAT16, " + "base tree decoding mask shall be INT8, output logits shall be FLOAT32."); + return false; + } + // Validate shapes of the input tensors. + bool const isBatchValid = baseTreeDecodingInputsEmbeds.getShape()[0] == mKVCache.getActiveBatchSize() + && baseTreeDecodingMask.getShape()[0] == mKVCache.getActiveBatchSize(); + if (!isBatchValid) + { + LOG_ERROR( + "eagleBaseTreeDecodingStepInputValidation(): Invalid batchSize of the input tensors. batchSize shall be " + "equal to the active batch " + "size set by the previous prefill stage."); + return false; + } + + int64_t const baseTreeDecodingSize = baseTreeDecodingInputsEmbeds.getShape()[1]; + bool const isBaseTreeDecodingSizeValid = baseTreeDecodingMask.getShape()[1] == baseTreeDecodingSize + && baseTreeDecodingMask.getShape()[2] == baseTreeDecodingSize + && baseTreeDecodingInputsEmbeds.getShape()[2] == mConfig.hiddenSize; + if (!isBaseTreeDecodingSizeValid) + { + LOG_ERROR( + "eagleBaseTreeDecodingStepInputValidation(): Invalid base tree decoding size of the input tensors. " + "Base tree decoding size %d, expected hiddenSize %d, current base tree decoding mask shape: %s, " + "inputsEmbeds shape: %s", + baseTreeDecodingSize, mConfig.hiddenSize, baseTreeDecodingMask.getShape().formatString().c_str(), + baseTreeDecodingInputsEmbeds.getShape().formatString().c_str()); + return false; + } + + bool const isOutputShapeValid = outputLogits.getShape()[0] == outputHiddenStates.getShape()[0] + && outputLogits.getShape()[1] == mConfig.outputVocabSize + && outputHiddenStates.getShape()[1] == mConfig.outputHiddenDim; + if (!isOutputShapeValid) + { + LOG_ERROR( + "eagleBaseTreeDecodingStepInputValidation(): Invalid shape of the output tensors. Logits shape shall be " + "[select-token-size, %d], hidden states shape shall be [select-token-size, %d], " + "current outputLogits shape: %s, outputHiddenStates shape: %s", + mConfig.outputVocabSize, mConfig.outputHiddenDim, outputLogits.getShape().formatString().c_str(), + outputHiddenStates.getShape().formatString().c_str()); + return false; + } + + return true; +} + +bool LLMEngineRunner::eagleBaseTreeDecodingStepPrepareInputs(rt::Tensor const& baseTreeDecodingInputsEmbeds, + rt::Tensor const& baseTreeDecodingMask, int32_t activeBatchSize, cudaStream_t stream) +{ + int32_t const baseTreeDecodingSize = static_cast(baseTreeDecodingInputsEmbeds.getShape()[1]); + int32_t const packedBaseTreeDecodingMaskLen = static_cast(divUp(baseTreeDecodingSize, 32)); + + // Prepare extra input for engine execution. Assemble mask, position indices, select token + // indices, sequence context lengths. + // We can obtain the sequence start index from KVCache, the current KVCache size denote the start index of the "next + // token" in the sequence. + rt::Tensor const& sequenceStartIndices = mKVCache.getKVCacheLengths(); + + // Prepare inputs for plugin-based attention + check::check(mSelectTokenIndices.reshape({activeBatchSize, baseTreeDecodingSize}), + "Tensor reshape failed"); // 2D tensor [batch, num_tokens] + check::check(mSequenceContextLengths.reshape({activeBatchSize}), "Tensor reshape failed"); + check::check(mEagleBasePositionIds.reshape({activeBatchSize, baseTreeDecodingSize}), "Tensor reshape failed"); + check::check(mEagleBasePackedMask.reshape({activeBatchSize, baseTreeDecodingSize, packedBaseTreeDecodingMaskLen}), + "Tensor reshape failed"); + + kernel::prepareEagleBaseTreeDecodingInputs(baseTreeDecodingMask, sequenceStartIndices, mEagleBasePackedMask, + mEagleBasePositionIds, mSelectTokenIndices, mSequenceContextLengths, stream); + return true; +} + +bool LLMEngineRunner::eagleBaseTreeDecodingStepBindTensors(rt::Tensor const& baseTreeDecodingInputsEmbeds, + rt::Tensor& outputLogits, rt::Tensor& outputHiddenStates, int32_t activeBatchSize) +{ + int32_t const baseTreeDecodingSize = static_cast(baseTreeDecodingInputsEmbeds.getShape()[1]); + // Bind the input and output tensor into the engine. RopeCosSinCache and KVCache are pre-bind during runner + // initialization. + bool setEngineIOStatus{true}; + setEngineIOStatus &= mTRTExecutionContext->setTensorAddress( + binding_names::kInputsEmbeds, const_cast(baseTreeDecodingInputsEmbeds.rawPointer())); + setEngineIOStatus &= mTRTExecutionContext->setInputShape( + binding_names::kInputsEmbeds, baseTreeDecodingInputsEmbeds.getShape().getTRTDims()); + setEngineIOStatus + &= mTRTExecutionContext->setTensorAddress(binding_names::kContextLengths, mSequenceContextLengths.rawPointer()); + setEngineIOStatus &= mTRTExecutionContext->setInputShape( + binding_names::kContextLengths, mSequenceContextLengths.getShape().getTRTDims()); + setEngineIOStatus + &= mTRTExecutionContext->setTensorAddress(binding_names::kLastTokenIds, mSelectTokenIndices.rawPointer()); + setEngineIOStatus &= mTRTExecutionContext->setInputShape( + binding_names::kLastTokenIds, mSelectTokenIndices.getShape().getTRTDims()); + setEngineIOStatus &= mTRTExecutionContext->setTensorAddress( + binding_names::kKVCacheStartIndex, mKVCache.getKVCacheLengths().rawPointer()); + setEngineIOStatus &= mTRTExecutionContext->setInputShape( + binding_names::kKVCacheStartIndex, mKVCache.getKVCacheLengths().getShape().getTRTDims()); + + // For MRope (VLM), reshape the RopeCosSinCache to match the activeBatchSize + if (mConfig.ropeConfig.type == RopeType::kMRope) + { + check::check(mPosEncCosSinCache.reshape({activeBatchSize, mConfig.maxKVCacheCapacity, mConfig.rotaryDim}), + "Tensor reshape failed"); + } + + setEngineIOStatus + &= mTRTExecutionContext->setInputShape(binding_names::kRopeCosSin, mPosEncCosSinCache.getShape().getTRTDims()); + + // Update KV cache shapes to match activeBatchSize + setEngineIOStatus &= this->bindKVCacheToEngine(activeBatchSize); + setEngineIOStatus &= this->bindConvStateToEngine(activeBatchSize); + setEngineIOStatus &= this->bindRecurrentStateToEngine(activeBatchSize); + + // Bind packed attention mask for plugin-based attention + setEngineIOStatus + &= mTRTExecutionContext->setTensorAddress(binding_names::kAttentionMask, mEagleBasePackedMask.rawPointer()); + setEngineIOStatus &= mTRTExecutionContext->setInputShape( + binding_names::kAttentionMask, mEagleBasePackedMask.getShape().getTRTDims()); + setEngineIOStatus + &= mTRTExecutionContext->setTensorAddress(binding_names::kAttentionPosId, mEagleBasePositionIds.rawPointer()); + setEngineIOStatus &= mTRTExecutionContext->setInputShape( + binding_names::kAttentionPosId, mEagleBasePositionIds.getShape().getTRTDims()); + + // Bind deepstack_embeds to dummy tensors for Qwen3VL models during Eagle base tree decoding + if (mConfig.numDeepstackFeatures > 0) + { + for (int32_t idx = 0; idx < mConfig.numDeepstackFeatures; ++idx) + { + std::string deepstackEmbedName = binding_names::formatDeepstackEmbedsName(idx); + setEngineIOStatus + &= mTRTExecutionContext->setTensorAddress(deepstackEmbedName.c_str(), mDummyInputTensor.rawPointer()); + setEngineIOStatus &= mTRTExecutionContext->setInputShape(deepstackEmbedName.c_str(), + rt::Coords{activeBatchSize, baseTreeDecodingSize, mConfig.hiddenSize}.getTRTDims()); + } + } + + // Bind the output tensor into the engine. + setEngineIOStatus &= mTRTExecutionContext->setTensorAddress(binding_names::kLogits, outputLogits.rawPointer()); + setEngineIOStatus + &= mTRTExecutionContext->setTensorAddress(binding_names::kOutputHiddenStates, outputHiddenStates.rawPointer()); + + if (!setEngineIOStatus) + { + LOG_ERROR("Failed to bind engine input and output tensors."); + return false; + } + + return true; +} + +bool LLMEngineRunner::executeEagleBaseTreeDecodingStep(rt::Tensor const& baseTreeDecodingInputsEmbeds, + rt::Tensor const& baseTreeDecodingMask, rt::Tensor& outputLogits, rt::Tensor& outputHiddenStates, + cudaStream_t stream) +{ + bool const validateInputStatus = this->eagleBaseTreeDecodingStepInputValidation( + baseTreeDecodingInputsEmbeds, baseTreeDecodingMask, outputLogits, outputHiddenStates); + if (!validateInputStatus) + { + LOG_ERROR("Eagle base tree decoding request not performed due to invalid input tensors."); + return false; + } + + int32_t const activeBatchSize = baseTreeDecodingInputsEmbeds.getShape()[0]; + + if (!eagleBaseTreeDecodingStepPrepareInputs( + baseTreeDecodingInputsEmbeds, baseTreeDecodingMask, activeBatchSize, stream)) + { + LOG_ERROR("Failed to prepare inputs."); + return false; + } + + // Launch cuda graph if available for this request, otherwise proceed with normal TensorRT engine execution step. + auto const graphHash + = baseKey(baseTreeDecodingInputsEmbeds, outputLogits, outputHiddenStates, mActiveLoraWeightsName); + if (mBaseTreeDecodingCudaGraphs.find(graphHash) != mBaseTreeDecodingCudaGraphs.end()) + { + LOG_DEBUG("Use pre-captured CUDA graph for eagle base tree decoding step."); + cudaGraphExec_t graphExec = mBaseTreeDecodingCudaGraphs[graphHash].second; + CUDA_CHECK(cudaGraphLaunch(graphExec, stream)); + } + else + { + bool setOptimizationProfileStatus{true}; + setOptimizationProfileStatus + &= mTRTExecutionContext->setOptimizationProfileAsync(kGENERATION_PROFILE_INDEX, stream); + if (!setOptimizationProfileStatus) + { + LOG_ERROR("Failed to set optimization profile to the engine"); + throw std::runtime_error("Failed to set optimization profile to the engine"); + } + + // Prepare and bind tensors using shared helper function + if (!eagleBaseTreeDecodingStepBindTensors( + baseTreeDecodingInputsEmbeds, outputLogits, outputHiddenStates, activeBatchSize)) + { + LOG_ERROR("Failed to bind tensors."); + return false; + } + + // launch the engine execution. + bool executeStatus{true}; + executeStatus &= mTRTExecutionContext->enqueueV3(stream); + if (!executeStatus) + { + LOG_ERROR("Failed on TensorRT eagle base tree decoding stage enqueueV3() call."); + return false; + } + } + + // Note in the base tree decoding step we explicitly don't commit the KVCache since we process the "whole tree" in + // these steps. + LOG_DEBUG("Eagle base tree decoding stage execution completed for request with batch size %d.", activeBatchSize); + return true; +} + +bool LLMEngineRunner::captureVanillaDecodingCudaGraph(rt::Tensor const& inputsEmbeds, rt::Tensor& outputLogits, + std::string const& loraWeightsPath, cudaStream_t stream, rt::OptionalOutputTensor outputHiddenStates, + rt::OptionalOutputTensor outputContextEmbeds) +{ + bool setOptimizationProfileStatus{true}; + setOptimizationProfileStatus + &= mTRTExecutionContext->setOptimizationProfileAsync(kGENERATION_PROFILE_INDEX, stream); + if (!setOptimizationProfileStatus) + { + LOG_ERROR("Failed to set optimization profile to the engine"); + throw std::runtime_error("Failed to set optimization profile to the engine"); + } + + auto const key = decodingKey(inputsEmbeds, outputLogits, loraWeightsPath); + if (mCudaGraphs.find(key) != mCudaGraphs.end()) + { + LOG_INFO("CUDA graph already captured for the input tensors with LoRA weights %s.", loraWeightsPath.c_str()); + return true; + } + + if (isLoraWeightsSupported() && !this->switchLoraWeights(loraWeightsPath)) + { + LOG_ERROR("Failed to switch LoRA weights to '%s', unable to capture CUDA graph.", loraWeightsPath.c_str()); + return false; + } + + // Here we will simulate the state of the EngineRunner after executing one prefill request for a batched request. + int32_t const activeBatchSize = inputsEmbeds.getShape()[0]; + constexpr int32_t simulateCacheLength{128}; + std::vector reuseKVCacheLengths(activeBatchSize, simulateCacheLength); + rt::Tensor const reuseKVCacheLengthsTensor(reuseKVCacheLengths.data(), {activeBatchSize}, rt::DeviceType::kCPU, + DataType::kINT32, "vanilla_reuse_kv_cache_lengths"); + + mKVCache.resetForNewSequences(reuseKVCacheLengthsTensor, stream); + + // Validate the input tensors. + bool const validateInputStatus = this->vanillaDecodingStepInputValidation(inputsEmbeds, outputLogits); + if (!validateInputStatus) + { + LOG_ERROR("Generation request is invalid, unable to capture CUDA graph."); + return false; + } + if (!vanillaDecodingStepPrepareInputs(activeBatchSize, stream)) + { + LOG_ERROR("Failed to prepare inputs."); + return false; + } + + if (!vanillaDecodingStepBindTensors( + inputsEmbeds, outputLogits, outputHiddenStates, outputContextEmbeds, activeBatchSize)) + { + LOG_ERROR("Failed to bind engine input and output tensors."); + return false; + } + + // launch the engine execution. This will trigger the shape machine of TensorRT engine to avoid cudaGraph capture + // error. + bool executeStatus{true}; + executeStatus &= mTRTExecutionContext->enqueueV3(stream); + if (!executeStatus) + { + LOG_ERROR("Failed on TensorRT engine enqueueV3() call."); + return false; + } + CUDA_CHECK(cudaStreamSynchronize(stream)); + + auto graphPair = captureTRTCudaGraph(mTRTExecutionContext.get(), stream); + if (!graphPair) + { + LOG_WARNING("Failed to capture CUDA graph."); + return false; + } + else + { + LOG_DEBUG( + "CUDA graph captured successfully for input shape %s with LoRA weights '%s' (Empty string if no LoRA " + "weights).", + inputsEmbeds.getShape().formatString().c_str(), loraWeightsPath.c_str()); + mCudaGraphs[key] = graphPair.value(); + return true; + } +} + +bool LLMEngineRunner::captureEagleBaseTreeDecodingCudaGraph(rt::Tensor const& baseTreeDecodingInputsEmbeds, + rt::Tensor const& baseTreeDecodingMask, rt::Tensor& outputLogits, rt::Tensor& outputHiddenStates, + std::string const& loraWeightsName, cudaStream_t stream) +{ + bool setOptimizationProfileStatus{true}; + setOptimizationProfileStatus + &= mTRTExecutionContext->setOptimizationProfileAsync(kGENERATION_PROFILE_INDEX, stream); + if (!setOptimizationProfileStatus) + { + LOG_ERROR("Failed to set optimization profile to the engine"); + throw std::runtime_error("Failed to set optimization profile to the engine"); + } + + auto const key = baseKey(baseTreeDecodingInputsEmbeds, outputLogits, outputHiddenStates, loraWeightsName); + if (mBaseTreeDecodingCudaGraphs.find(key) != mBaseTreeDecodingCudaGraphs.end()) + { + LOG_INFO("CUDA graph already captured for the input tensors with LoRA weights %s.", loraWeightsName.c_str()); + return true; + } + + if (isLoraWeightsSupported() && !this->switchLoraWeights(loraWeightsName)) + { + LOG_ERROR("Failed to switch LoRA weights to '%s', unable to capture CUDA graph.", loraWeightsName.c_str()); + return false; + } + + // Here we will simulate the state of the EngineRunner after executing one prefill request for a batched request. + int32_t const activeBatchSize = baseTreeDecodingInputsEmbeds.getShape()[0]; + constexpr int32_t simulateCacheLength{128}; + std::vector reuseKVCacheLengths(activeBatchSize, simulateCacheLength); + rt::Tensor const reuseKVCacheLengthsTensor( + reuseKVCacheLengths.data(), {activeBatchSize}, rt::DeviceType::kCPU, DataType::kINT32); + + mKVCache.resetForNewSequences(reuseKVCacheLengthsTensor, stream); + + bool const validateInputStatus = this->eagleBaseTreeDecodingStepInputValidation( + baseTreeDecodingInputsEmbeds, baseTreeDecodingMask, outputLogits, outputHiddenStates); + if (!validateInputStatus) + { + LOG_ERROR("Eagle base tree decoding request not performed due to invalid input tensors."); + return false; + } + + // Prepare and bind tensors using shared helper function + if (!eagleBaseTreeDecodingStepPrepareInputs( + baseTreeDecodingInputsEmbeds, baseTreeDecodingMask, activeBatchSize, stream)) + { + LOG_ERROR("Failed to prepare inputs."); + return false; + } + + if (!eagleBaseTreeDecodingStepBindTensors( + baseTreeDecodingInputsEmbeds, outputLogits, outputHiddenStates, activeBatchSize)) + { + LOG_ERROR("Failed to bind tensors."); + return false; + } + + // launch the engine execution. This will trigger the shape machine of TensorRT engine to avoid cudaGraph capture. + // error. + bool executeStatus{true}; + executeStatus &= mTRTExecutionContext->enqueueV3(stream); + + if (!executeStatus) + { + LOG_ERROR("Failed on TensorRT eagle base tree decoding stage enqueueV3() call."); + return false; + } + + CUDA_CHECK(cudaStreamSynchronize(stream)); + + auto graphPair = captureTRTCudaGraph(mTRTExecutionContext.get(), stream); + if (!graphPair) + { + LOG_WARNING("Failed to capture CUDA graph."); + return false; + } + else + { + LOG_DEBUG( + "CUDA graph captured successfully for input shape %s with LoRA weights '%s' (Empty string if no LoRA " + "weights).", + baseTreeDecodingInputsEmbeds.getShape().formatString().c_str(), loraWeightsName.c_str()); + mBaseTreeDecodingCudaGraphs[key] = graphPair.value(); + return true; + } +} + +bool LLMEngineRunner::resetLoraWeights() +{ + if (!isLoraWeightsSupported()) + { + return true; + } + mActiveLoraWeightsName = ""; + bool resetStatus{true}; + for (auto const& loraWeightsTensorName : getLoraWeightsTensorNames()) + { + nvinfer1::Dims emptyLoraShape + = mEngine->getProfileShape(loraWeightsTensorName.c_str(), 0, nvinfer1::OptProfileSelector::kMAX); + + // Use dummy tensor as zero tensor for LoRA weights + resetStatus + &= mTRTExecutionContext->setTensorAddress(loraWeightsTensorName.c_str(), mDummyInputTensor.rawPointer()); + + // Set shape to kEMPTY_LORA_RANK and assign zero value tensor to disable LoRA + if (loraWeightsTensorName.find(binding_names::kLoraAPrefix) != std::string::npos) + { + // LoRA A has shape [k, rank], set rank to kEMPTY_LORA_RANK + emptyLoraShape.d[1] = kEMPTY_LORA_RANK; + } + else if (loraWeightsTensorName.find(binding_names::kLoraBPrefix) != std::string::npos) + { + // LoRA B has shape [rank, n], set rank to kEMPTY_LORA_RANK + emptyLoraShape.d[0] = kEMPTY_LORA_RANK; + } + resetStatus &= mTRTExecutionContext->setInputShape(loraWeightsTensorName.c_str(), emptyLoraShape); + if (!resetStatus) + { + LOG_ERROR("Failed to reset LoRA weights: %s", loraWeightsTensorName.c_str()); + return false; + } + } + return resetStatus; +} + +bool LLMEngineRunner::addLoraWeights( + std::string const& loraWeightsName, std::string const& loraWeightsPath, cudaStream_t stream) +{ + if (!isLoraWeightsSupported()) + { + LOG_ERROR("addLoraWeights(): Engine does not support LoRA weights."); + } + + if (mLoraWeights.find(loraWeightsName) != mLoraWeights.end()) + { + LOG_ERROR("addLoraWeights(): LoRA weights %s already added", loraWeightsName.c_str()); + return false; + } + + // Load tensors using the new unified interface + std::vector tensors; + if (!safetensors::loadSafetensors(loraWeightsPath, tensors, stream)) + { + LOG_ERROR("addLoraWeights(): Failed to load LoRA weights %s from: %s", loraWeightsName.c_str(), + loraWeightsPath.c_str()); + return false; + } + + // Validate the LoRA weights do not exceed the max LoRA rank + for (auto const& tensor : tensors) + { + if (tensor.getName().find(binding_names::kLoraAPrefix) != std::string::npos) + { + if (tensor.getShape()[1] > mConfig.maxSupportedLoraRank) + { + LOG_ERROR("addLoraWeights(): LoRA A (%s) tensor's rank (%d) exceeds the max LoRA rank (%d)", + tensor.getName().c_str(), tensor.getShape()[1], mConfig.maxSupportedLoraRank); + return false; + } + } + else if (tensor.getName().find(binding_names::kLoraBPrefix) != std::string::npos) + { + if (tensor.getShape()[0] > mConfig.maxSupportedLoraRank) + { + LOG_ERROR("addLoraWeights(): LoRA B (%s) tensor's rank (%d) exceeds the max LoRA rank (%d)", + tensor.getName().c_str(), tensor.getShape()[0], mConfig.maxSupportedLoraRank); + return false; + } + } + } + + // Store the tensors in our map + mLoraWeights[loraWeightsName] = std::move(tensors); + LOG_INFO("addLoraWeights(): Added LoRA weights %s from: %s", loraWeightsName.c_str(), loraWeightsPath.c_str()); + return true; +} + +std::vector LLMEngineRunner::getLoraWeightsTensorNames() const +{ + std::vector loraWeightsTensorNames; + // Get the number of bindings in the engine + int32_t numBindings = mEngine->getNbIOTensors(); + for (int32_t i = 0; i < numBindings; ++i) + { + char const* bindingName = mEngine->getIOTensorName(i); + std::string bindingNameStr(bindingName); + if (binding_names::isLoraBinding(bindingNameStr)) + { + loraWeightsTensorNames.push_back(bindingNameStr); + } + } + return loraWeightsTensorNames; +} + +bool LLMEngineRunner::switchLoraWeights(std::string const& loraWeightsName) +{ + if (!isLoraWeightsSupported()) + { + LOG_ERROR("switchLoraWeights(): API call is invalid. LLM engine does not support LoRA weights."); + return false; + } + if (loraWeightsName.empty()) + { + this->resetLoraWeights(); + LOG_DEBUG("switchLoraWeights(): Switched to no LoRA weights."); + return true; + } + + // Check if the requested LoRA exists + auto it = mLoraWeights.find(loraWeightsName); + if (it == mLoraWeights.end()) + { + LOG_ERROR("switchLoraWeights(): LoRA weights with name '%s' not found", loraWeightsName.c_str()); + return false; + } + + auto& loraTensors = it->second; + + // Iterate through all LoRA weights bindings + for (auto const& loraWeightsTensorName : this->getLoraWeightsTensorNames()) + { + // Try to find the tensor in the LoRA weights + auto loraTensorIt = std::find_if(loraTensors.begin(), loraTensors.end(), + [loraWeightsTensorName](rt::Tensor const& tensor) { return tensor.getName() == loraWeightsTensorName; }); + + bool setLoraWeightsStatus{true}; + + if (loraTensorIt != loraTensors.end()) + { + // Found matching tensor, use its data + setLoraWeightsStatus + &= mTRTExecutionContext->setInputShape(loraWeightsTensorName.c_str(), loraTensorIt->getTRTDims()); + setLoraWeightsStatus + &= mTRTExecutionContext->setInputShape(loraWeightsTensorName.c_str(), loraTensorIt->getTRTDims()); + setLoraWeightsStatus + &= mTRTExecutionContext->setTensorAddress(loraWeightsTensorName.c_str(), loraTensorIt->rawPointer()); + setLoraWeightsStatus + &= mTRTExecutionContext->setTensorAddress(loraWeightsTensorName.c_str(), loraTensorIt->rawPointer()); + LOG_DEBUG("switchLoraWeights(): LoRA weights tensor with name '%s' found. Set shape to %s.", + loraWeightsTensorName.c_str(), loraTensorIt->getShape().formatString().c_str()); + } + else + { + // Tensor not found in this LoRA adapter, use dummy tensor as zero tensor with shape kEMPTY_LORA_RANK + nvinfer1::Dims shape + = mEngine->getProfileShape(loraWeightsTensorName.c_str(), 0, nvinfer1::OptProfileSelector::kMAX); + if (loraWeightsTensorName.find(binding_names::kLoraAPrefix) != std::string::npos) + { + // LoRA A has shape [k, rank], set rank to kEMPTY_LORA_RANK + shape.d[1] = kEMPTY_LORA_RANK; + } + else if (loraWeightsTensorName.find(binding_names::kLoraBPrefix) != std::string::npos) + { + // LoRA B has shape [rank, n], set rank to kEMPTY_LORA_RANK + shape.d[0] = kEMPTY_LORA_RANK; + } + setLoraWeightsStatus &= mTRTExecutionContext->setInputShape(loraWeightsTensorName.c_str(), shape); + setLoraWeightsStatus &= mTRTExecutionContext->setTensorAddress( + loraWeightsTensorName.c_str(), mDummyInputTensor.rawPointer()); + LOG_DEBUG( + "LoRA weights tensor with name '%s' not found. Set shape to rank %d with zero " + "tensor.", + loraWeightsTensorName.c_str(), kEMPTY_LORA_RANK); + } + if (!setLoraWeightsStatus) + { + LOG_ERROR("Failed to set LoRA weights: %s", loraWeightsTensorName.c_str()); + return false; + } + } + // Set the active LoRA weights name + mActiveLoraWeightsName = loraWeightsName; + LOG_DEBUG("switchLoraWeights(): Switched to LoRA weights with name '%s'.", loraWeightsName.c_str()); + return true; +} + +std::string LLMEngineRunner::getActiveLoraWeightsName() const +{ + return mActiveLoraWeightsName; +} + +std::vector LLMEngineRunner::getAvailableLoraWeights() const +{ + std::vector loraWeightsNames; + for (auto const& [loraWeightsName, _] : mLoraWeights) + { + loraWeightsNames.push_back(loraWeightsName); + } + return loraWeightsNames; +} + +bool LLMEngineRunner::isLoraWeightsSupported() const noexcept +{ + return mConfig.maxSupportedLoraRank > 0; +} + +int32_t LLMEngineRunner::getMaxLoraWeightsDimension() const +{ + if (!isLoraWeightsSupported()) + { + return 0; + } + + int32_t maxDim = 0; + + // Query engine profile shapes for all LoRA weight tensors + for (auto const& loraWeightsTensorName : getLoraWeightsTensorNames()) + { + nvinfer1::Dims maxShape + = mEngine->getProfileShape(loraWeightsTensorName.c_str(), 0, nvinfer1::OptProfileSelector::kMAX); + + if (loraWeightsTensorName.find(binding_names::kLoraAPrefix) != std::string::npos) + { + // LoRA A has shape [k, rank], we want max k + maxDim = std::max(maxDim, static_cast(maxShape.d[0])); + } + else if (loraWeightsTensorName.find(binding_names::kLoraBPrefix) != std::string::npos) + { + // LoRA B has shape [rank, n], we want max n + maxDim = std::max(maxDim, static_cast(maxShape.d[1])); + } + } + + return maxDim; +} + +} // namespace rt +} // namespace trt_edgellm diff --git a/cpp/runtime/llmEngineRunner.h b/cpp/runtime/llmEngineRunner.h new file mode 100644 index 00000000..babb61d4 --- /dev/null +++ b/cpp/runtime/llmEngineRunner.h @@ -0,0 +1,441 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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. + */ + +#pragma once + +#include "common/hashUtils.h" +#include "common/tensor.h" +#include "runtime/linearKVCache.h" +#include "runtime/llmRuntimeUtils.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace trt_edgellm +{ +namespace rt +{ +using Json = nlohmann::json; + +/*! + * @brief Configuration structure for LLM engine runner + * + * Contains all runtime configuration parameters for the LLM engine. + */ +struct LLMEngineRunnerConfig +{ + RopeConfig ropeConfig{}; //!< Type of rotary positional encoding + bool useContextDependentRope{false}; //!< Use context-dependent RoPE + bool enableEagleSpecDecode{false}; //!< Enable Eagle speculative decoding + bool enableContextEmb{false}; //!< GR00T VLA: engine also outputs context_embs + bool enableLmHiddenStates{false}; //!< GR00T VLA: engine outputs lm_hidden_states for action_context engine + bool enablePrefixKVOutputs{false}; //!< VLA: engine outputs prefix_k / prefix_v for action head + bool hasLogitsOutput{true}; //!< False for prefill-only VLA language engines without lm_head + bool hasLastTokenIdsInput{true}; //!< False when the exported language engine omits last_token_ids + bool useTrtNativeOps{false}; //!< Use TensorRT native operations instead of custom plugin + int32_t numDecoderLayers{}; //!< Number of decoder layers + int32_t numKVHeads{}; //!< Number of key-value heads + int32_t headDim{}; //!< Dimension of each attention head + int32_t rotaryDim{}; //!< Rotary embedding dimension + int32_t hiddenSize{}; //!< Model's hidden dimension + int32_t maxSupportedBatchSize{}; //!< Maximum supported batch size + int32_t maxSupportedInputLength{}; //!< Maximum supported input length + int32_t maxKVCacheCapacity{}; //!< Maximum KV cache capacity + int32_t vocabSize{}; //!< Vocabulary size (full vocabulary) + int32_t reducedVocabSize{0}; //!< Reduced vocabulary size (0 if not using reduced vocab) + int32_t outputVocabSize{}; //!< Actual output vocabulary size (reducedVocabSize if enabled, else vocabSize) + int32_t maxSupportedLoraRank{}; //!< Maximum supported LoRA rank + int32_t outputHiddenDim{}; //!< Output hidden dimension for Eagle speculative decoding (hidden_size * 3) + int32_t maxVerifyTreeSize{}; //!< Maximum verification tree size for Eagle speculative decoding + int32_t contextEmbDim{}; //!< context_embs hidden dim for GR00T VLA action head + std::vector prefixKVOutputShape{}; //!< PI0.5 prefix_k / prefix_v max output shape from engine + int32_t numDeepstackFeatures{0}; //!< Number of deepstack features for Qwen3-VL and Qwen3-Omni + int32_t audioTokenId{0}; //!< Special token ID for audio in Qwen3-Omni + int32_t imageTokenId{0}; //!< Special token ID for image in Qwen3-Omni + + // Hybrid model configuration + int32_t numLinearAttnLayers{0}; //!< Number of recurrent layers (0 for pure attention models) + int32_t numAttentionLayers{0}; //!< Number of attention layers (equals numDecoderLayers for pure attention) + int32_t recurrentStateNumHeads{0}; //!< Number of recurrent heads (hv for GDN, mamba_num_heads for Mamba) + int32_t recurrentStateHeadDim{0}; //!< Dimension of each recurrent head (k for GDN, mamba_head_dim for Mamba) + int32_t recurrentStateSize{0}; //!< Recurrent state dimension (v for GDN, dstate for Mamba) + int32_t convDim{0}; //!< Conv1d channel dimension + int32_t convKernel{0}; //!< Conv1d kernel width +}; + +//! The class wraps the TensorRT engine built for auto-regressive style decoder model. +//! The LLMEngineRunner define the interface for upper level runtime to execute engine actions to drive +//! autoregressive decoding with/without speculative decoding for edge inference scenarios. Current design +//! assume prefill and decoding operations are synchronous so a batched requests need to perform prefill and +//! decoding at the same time (no continuous batching). +//! The LLMEngineRunner will: +//! 1. Hold TensorRT resources of the LLM engine (TRT IRuntime, CUDA Engine, Execution Contexts). +//! 2. Hold the LinearKVCache resources that support till maxSupportedBatchSize and maxSequenceLength. +//! 3. Hold the Rope CosSinCache tensor required for positional encoding. +class LLMEngineRunner +{ +public: + /*! + * @brief Construct LLM engine runner + * @param enginePath Path to TensorRT engine file + * @param configPath Path to model configuration file + * @param loraWeightsMap Map of LoRA weight names to file paths + * @param stream CUDA stream for operations + * @throws std::runtime_error If engine loading, configuration parsing, or initialization fails, or a CUDA operation + * fails + */ + LLMEngineRunner(std::filesystem::path const& enginePath, std::filesystem::path const& configPath, + std::unordered_map const& loraWeightsMap, cudaStream_t stream); + + //! @brief Destructor + ~LLMEngineRunner() noexcept; + + /*! + * @brief Get the required context memory size for this engine + * @return Required context memory size in bytes + */ + int64_t getRequiredContextMemorySize() const; + + /*! + * @brief Set shared context memory for the execution context + * @param sharedContextMemory Tensor containing the shared device memory (must be on GPU) + * @return True on success, false if the tensor is too small + * @note The tensor size must be >= getRequiredContextMemorySize(). Must be called before execution. + */ + bool setContextMemory(rt::Tensor& sharedContextMemory); + + //! API entry to get the Rope CosSinCache tensor. + //! The API is useful when the rope cos/sin cache depends on the context which cannot be initialized + //! in advance when creating the LLMEngineRunner instance. + rt::Tensor& getRopeCosSinCacheTensor() noexcept; + + //! @brief Get reference to the linear KV cache (also owns recurrent/conv state buffers for hybrid models) + //! @return Reference to LinearKVCache + rt::LinearKVCache& getLinearKVCache() noexcept; + + //! @brief Get engine configuration + //! @return Engine configuration structure + LLMEngineRunnerConfig getEngineConfig() const noexcept; + + //! @brief Set an extra input tensor for the engine + //! + //! This is a temporary API for binding additional input tensors that are not part of + //! the standard LLM input set. + //! @note This is not a good design but we put it here temporarily to support TTS inference. + //! @note The API will be replaced soon with a better design. Please don't follow this schema. + //! + //! Example use case: CodePredictor's lm_head_weight input for dynamic lm_head selection. + //! + //! @param name The name of the LMHead input weights in the ONNX/TRT model + //! @param tensor The tensor to bind (must be on GPU, shape must match engine expectation) + //! @return True if the binding was successful + //! @note Must be called before executePrefillStep/executeVanillaDecodingStep + bool setLMHeadWeights(std::string const& name, rt::Tensor const& tensor); + + //! API entry to execute one prefill engine action for a batched request. The API will clear existing KVCache for + //! last + //! batch of requests and perform prefill operations to fill the KVCache and produce the output logits. + //! Inputs: + //! inputsEmbeds [GPU]: The input embeddings for the batch of new requests [batchSize, seqLen, hiddenSize]. + //! contextLengths [CPU]: The context lengths for each sequence in the batch. + //! deepstackEmbeds [GPU]: Optional. Deepstack embeddings for Qwen3-VL (already embedded). + //! outputLogits [GPU]: The output logits for the batch of requests. + //! outputHiddenStates [GPU]: Optional. The output hidden states for Eagle speculative decoding. + //! outputContextEmbeds [GPU]: Optional. The output context_embs for GR00T VLA action head. + //! stream: The CUDA stream to execute the prefill step. + //! Returns: + //! True if the prefill step is successful, false otherwise. + //! @throws std::runtime_error if setting optimization profile fails, or a CUDA operation fails + bool executePrefillStep(rt::Tensor const& inputsEmbeds, rt::Tensor const& contextLengths, + rt::OptionalInputTensors deepstackEmbeds, rt::Tensor& outputLogits, rt::OptionalOutputTensor outputHiddenStates, + cudaStream_t stream, rt::OptionalOutputTensor outputContextEmbeds = std::nullopt, + rt::OptionalOutputTensor outputPrefixK = std::nullopt, rt::OptionalOutputTensor outputPrefixV = std::nullopt); + + //! API entry to execute one vanilla decoding engine action for a batched request. The API will perform decoding + //! operations fill the KVCache of the new generated tokens and produce the output logits. The decoding + //! operation shall be performed after the prefill step is completed. + //! Inputs: + //! inputsEmbeds [GPU]: The input embeddings for the batch of new requests [batchSize, 1, hiddenSize]. + //! stream: The CUDA stream to execute the decoding step. + //! Outputs: + //! outputLogits [GPU]: The output logits for the batch of requests. + //! Returns: + //! True if the decoding step is successful, false otherwise. + //! @throws std::runtime_error if setting optimization profile fails, or a CUDA operation fails + bool executeVanillaDecodingStep(rt::Tensor const& inputsEmbeds, rt::Tensor& outputLogits, + rt::OptionalOutputTensor outputHiddenStates, cudaStream_t stream, + rt::OptionalOutputTensor outputContextEmbeds = std::nullopt); + + //! API entry to execute eagle base tree decoding step. The API will takes a draft tree of input embeddings. + //! baseTreeDecodingMask denote the relationship between the draft tree nodes. + //! Inputs: + //! baseTreeDecodingInputsEmbeds [GPU, Float16]: Input embeddings for the base model with shape [batchSize, + //! Tree-Size, hiddenSize]. baseTreeDecodingMask [GPU, Int32]: Denote the relationship between the base tree + //! nodes with shape + //! [batchSize, Tree-Size, Tree-Size]. + //! stream: The CUDA stream to execute the base tree decoding step. + //! Outputs: + //! outputLogits [GPU, Float16]: The output logits with shape [batchSize*Tree-Size, base-Vocab-Size]. + //! outputHiddenStates [GPU]: The output hidden states with shape [batchSize*Tree-Size, base-hidden-dim]. + //! @throws std::runtime_error if setting optimization profile fails, or a CUDA operation fails + bool executeEagleBaseTreeDecodingStep(rt::Tensor const& baseTreeDecodingInputsEmbeds, + rt::Tensor const& baseTreeDecodingMask, rt::Tensor& outputLogits, rt::Tensor& outputHiddenStates, + cudaStream_t stream); + + //! API entry to capture the CUDA graph for the decoding step. If CUDA graph capture is successful, later + //! call to executeVanillaDecodingStep() will always launch the captured CUDA graph. + //! Inputs: + //! inputsEmbeds [GPU]: The input embeddings for the batch of new requests [batchSize, 1, hiddenSize]. + //! outputLogits [GPU]: The output logits for the batch of requests. + //! loraWeightsName: The name to the LoRA weights. Empty string if no LoRA weights. + //! stream: The CUDA stream to execute the decoding step. + //! Returns: + //! True if the CUDA graph capture is successful, false otherwise. + //! @throws std::runtime_error if setting optimization profile fails, or a CUDA operation fails + bool captureVanillaDecodingCudaGraph(rt::Tensor const& inputsEmbeds, rt::Tensor& outputLogits, + std::string const& loraWeightsName, cudaStream_t stream, + rt::OptionalOutputTensor outputHiddenStates = std::nullopt, + rt::OptionalOutputTensor outputContextEmbeds = std::nullopt); + + //! API entry to switch the LoRA weights of the LLM engine. + //! Inputs: + //! loraWeightsName: The name of the LoRA weights. + //! Returns: + //! True if the LoRA weights switch is successful, false otherwise. + bool switchLoraWeights(std::string const& loraWeightsName); + + //! API entry to get the active LoRA weights name. + //! Returns: + //! The active LoRA weights name. + std::string getActiveLoraWeightsName() const; + + //! API entry to get the LoRA weights. + //! Returns: + //! The LoRA weights names. + std::vector getAvailableLoraWeights() const; + + //! API entry to capture the CUDA graph for the base model tree decoding step. If CUDA graph capture is successful, + //! later + //! call to executeEagleBaseTreeDecodingStep() will always launch the captured CUDA graph. + //! Inputs: + //! baseTreeDecodingInputsEmbeds [GPU, Float16]: Input embeddings for the base model with shape [batchSize, + //! Tree-Size, hiddenSize]. baseTreeDecodingMask [GPU, Int32]: Denote the relationship between the base tree + //! nodes with shape + //! [batchSize, Tree-Size, Tree-Size]. + //! outputLogits [GPU, Float16]: The output logits with shape [batchSize*Tree-Size, base-Vocab-Size]. + //! outputHiddenStates [GPU]: The output hidden states with shape [batchSize*Tree-Size, base-hidden-dim]. + //! stream: The CUDA stream to capture the CUDA graph. The API will capture the CUDA graph for the base tree + //! decoding step. + //! Returns: + //! True if the CUDA graph capture is successful, false otherwise. + //! @throws std::runtime_error if setting optimization profile fails, or a CUDA operation fails + bool captureEagleBaseTreeDecodingCudaGraph(rt::Tensor const& baseTreeDecodingInputsEmbeds, + rt::Tensor const& baseTreeDecodingMask, rt::Tensor& outputLogits, rt::Tensor& outputHiddenStates, + std::string const& loraWeightsName, cudaStream_t stream); + + //! Key to uniquely identify a captured CUDA graph for the decoding step + using DecodingGraphKey = std::tuple; + + //! Key to uniquely identify a captured CUDA graph for the base model verification step + using BaseGraphKey = std::tuple; + +private: + std::unique_ptr mRuntime; //!< TensorRT runtime + std::unique_ptr mEngine; //!< TensorRT engine + std::unique_ptr mTRTExecutionContext; //!< Prefill and Generation execution context + + //! Holds the CUDA graph captured for the decoding step. Each CUDA graph is associated with a unique key value + //! which denote the input/output shapes and other execution properties like LoRA weights. + hash_utils::HashMap> mCudaGraphs; + + //! Holds the CUDA graph captured for the base model verification step. Each CUDA graph is associated with a unique + //! key value which denote the input/output shapes and other execution properties. + hash_utils::HashMap> mBaseTreeDecodingCudaGraphs; + + //! Holds the LoRA weights for the LLM engine. + std::unordered_map> mLoraWeights{}; + std::string mActiveLoraWeightsName{}; //!< Name of currently active LoRA weights + + LLMEngineRunnerConfig mConfig{}; //!< Engine configuration + + //! The Rope CosSinCache tensor that pre-computed prior to engine execution. + //! The design is to produce better performance and accommodate complex context dependent rope. + rt::Tensor mPosEncCosSinCache{}; + + //! The select token indices tensor is used to select indices from hidden states to pass to + //! the LM head of LLM model. Enforce to be int64_t to align with ONNX Gather-ND specification. + rt::Tensor mSelectTokenIndices{}; + rt::Tensor mHostSelectTokenIndices{}; //!< Host tensor for select token indices (pinned memory) + + //! The tensor has different meaning for prefill and decoding phase due to implementation of + //! the AttentionPlugin. Used as LLM engine input. + //! For prefill phase, the field denotes the actual content length of input_ids for each sequence. + //! For decoding phase, this field denotes the cumulative length of the sequence length of prefill + //! plus generated tokens (including the length in "current" run). + rt::Tensor mSequenceContextLengths{}; + + //! The LinearKVCache tensor that carried for the LLM model execution. + //! Also owns recurrent and conv state buffers for hybrid models. + rt::LinearKVCache mKVCache{}; + + //! Dummy input tensor used to reserve space for unused input tensors. We always keep this tensor as zero tensor + //! because to "void" some computation (ex. use as empty lora weights as if there is no LoRA GEMM). + rt::Tensor mDummyInputTensor{}; + + //! Dummy output tensor used to reserve space for unused output tensors. TRT engines have static I/O, to keep + //! runtime design clean, we will route unused output tensors to this dummy tensor. + rt::Tensor mDummyOutputTensor{}; + + //! Dummy output tensor for context_embs / lm_hidden_states when the engine exports it but the caller does not need + //! it. + rt::Tensor mDummyContextEmbTensor{}; + + //! Prefill auxiliary sequence output binding (context_embs or lm_hidden_states). + std::string mPrefillAuxOutputName{}; + + //! The eagle base position ids tensor within the sequence that used by positional encoding. + rt::Tensor mEagleBasePositionIds{}; + + //! The eagle base packed mask tensor to indicate the attention relationship between the base verify nodes. + rt::Tensor mEagleBasePackedMask{}; + + /*! + * @brief Initialize configuration from JSON file + * @param configJson JSON configuration object + * @return True on success, false on failure + */ + bool initializeConfigFromJson(Json const& configJson) noexcept; + + /*! + * @brief Validate configuration against engine + * @return True if valid, false otherwise + */ + bool validateConfigFromEngine(); + + /*! + * @brief Bind KV cache to engine for prefill and generation of new requests + * @param activeBatchSize Number of active sequences + * @return True on success, false on failure + */ + bool bindKVCacheToEngine(int32_t activeBatchSize); + + //! @brief Validate inputs for prefill step + bool prefillStepInputValidation(rt::Tensor const& inputsEmbeds, rt::Tensor const& contextLengths, + rt::Tensor const& outputLogits, rt::OptionalOutputTensor outputHiddenStates, + rt::OptionalOutputTensor outputContextEmbeds, rt::OptionalInputTensors deepstackEmbeds) noexcept; + + //! @brief Validate inputs for vanilla decoding step + bool vanillaDecodingStepInputValidation(rt::Tensor const& inputsEmbeds, rt::Tensor const& outputLogits) noexcept; + + //! @brief Prepare inputs for vanilla decoding step (shared between execute and capture) + bool vanillaDecodingStepPrepareInputs(int32_t activeBatchSize, cudaStream_t stream); + + //! @brief Bind tensors for vanilla decoding step (shared between execute and capture) + bool vanillaDecodingStepBindTensors(rt::Tensor const& inputsEmbeds, rt::Tensor& outputLogits, + rt::OptionalOutputTensor outputHiddenStates, rt::OptionalOutputTensor outputContextEmbeds, + int32_t activeBatchSize); + + //! @brief Validate inputs for Eagle base tree decoding step + bool eagleBaseTreeDecodingStepInputValidation(rt::Tensor const& baseTreeDecodingInputsEmbeds, + rt::Tensor const& baseTreeDecodingMask, rt::Tensor const& outputLogits, + rt::Tensor const& outputHiddenStates) noexcept; + + //! @brief Prepare and bind tensors for Eagle base tree decoding step (shared between execute and capture) + bool eagleBaseTreeDecodingStepBindTensors(rt::Tensor const& baseTreeDecodingInputsEmbeds, rt::Tensor& outputLogits, + rt::Tensor& outputHiddenStates, int32_t activeBatchSize); + + bool eagleBaseTreeDecodingStepPrepareInputs(rt::Tensor const& baseTreeDecodingInputsEmbeds, + rt::Tensor const& baseTreeDecodingMask, int32_t activeBatchSize, cudaStream_t stream); + + //! The Function is used to add a LoRA weights to the LLM engine. + bool addLoraWeights(std::string const& loraWeightsName, std::string const& loraWeightsPath, cudaStream_t stream); + + /*! + * @brief Reset LoRA weights to dummy tensors with rank 0 + * @return True on success, false on failure + */ + bool resetLoraWeights(); + + /*! + * @brief Get maximum dimension required for LoRA weights across all LoRA bindings + * @return Maximum dimension (k for LoRA A, n for LoRA B), or 0 if no LoRA bindings + */ + int32_t getMaxLoraWeightsDimension() const; + + /*! + * @brief Get tensor names of LoRA weights + * @return Vector of LoRA weight tensor names + */ + std::vector getLoraWeightsTensorNames() const; + + //! @brief Check if LoRA weights are supported + //! @return True if supported, false otherwise + bool isLoraWeightsSupported() const noexcept; + + //! @brief Get the KV cache type + //! @return The KV cache type + nvinfer1::DataType getKVCacheType() const; + + //! @brief Get the recurrent state dtype from the engine binding (layer 0) + nvinfer1::DataType getRecurrentStateType() const; + + //! @brief Get the conv state dtype from the engine binding (layer 0) + nvinfer1::DataType getConvStateType() const; + + //! @brief Validate the KV cache type consistency + //! @return True if the KV cache type is consistent, false otherwise + //! @throws std::runtime_error if KV cache has mismatching data type + bool validateKVCacheType() const; + +private: + /*! + * @brief Bind KV cache to engine for prefill and generation of new requests (plugin path) + * @param activeBatchSize Number of active sequences + * @return True on success, false on failure + */ + bool bindPluginKVCacheToEngine(int32_t activeBatchSize); + + /*! + * @brief Bind separate K and V caches to engine for new requests (TRT native path) + * @param activeBatchSize Number of active sequences + * @return True on success, false on failure + */ + bool bindTRTNativeKVCacheToEngine(int32_t activeBatchSize); + + /*! + * @brief Bind recurrent state buffers for recurrent layers to the engine + * @param activeBatchSize Number of active sequences + * @return True on success, false on failure + */ + bool bindRecurrentStateToEngine(int32_t activeBatchSize); + + /*! + * @brief Bind conv state tensors to the TensorRT execution context + * + * @param activeBatchSize Current batch size to bind + * @return True on success, false on failure + */ + bool bindConvStateToEngine(int32_t activeBatchSize); +}; + +} // namespace rt +} // namespace trt_edgellm diff --git a/cpp/runtime/llmRuntimeUtils.h b/cpp/runtime/llmRuntimeUtils.h index 632f44d7..9d5b31a0 100644 --- a/cpp/runtime/llmRuntimeUtils.h +++ b/cpp/runtime/llmRuntimeUtils.h @@ -101,6 +101,10 @@ struct LLMGenerationRequest std::optional> pastTrajectory; //!< Optional past trajectory for Alpamayo (e.g. ego x,y,z history) + //!< Optional packed robot state for GR00T diffusion (row-major, length = state_horizon * state_dim) + std::vector robotState; + std::optional embodimentId; //!< Optional embodiment index for GR00T action head + //! Stop strings; generation halts on the earliest match and trims it from output. std::vector stopStrings; @@ -145,6 +149,11 @@ struct LLMGenerationRequest //! Called after cudaStreamSynchronize inside the decode loop. //! When nullopt (default), zero overhead — no callback is invoked. std::optional onTokenGenerated; + + // Number of trajectories the action stage should produce. + int32_t actionBatchSize{0}; + //!< Optional default embodiment index for GR00T action head (overrides export default when set). + std::optional embodimentId; }; /*! \brief LLM Generation Response structure @@ -158,6 +167,9 @@ struct LLMGenerationResponse std::vector outputAudios; //!< Generated audio data (Qwen3-Omni only) + //!< Denoised robot actions per batch item (row-major [action_horizon * action_dim]); populated by the VLA runtime + std::vector> outputActions; + //! Why each request halted (EOS, length, stop string, cancel, error); see `runtime/streaming.h`. std::vector finishReasons; }; diff --git a/cpp/runtime/motBackboneRunner.cpp b/cpp/runtime/motBackboneRunner.cpp new file mode 100644 index 00000000..704fbac6 --- /dev/null +++ b/cpp/runtime/motBackboneRunner.cpp @@ -0,0 +1,495 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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. + */ + +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/motBackboneRunner.h" + +#include "common/checkMacros.h" +#include "common/cudaUtils.h" +#include "common/logger.h" +#include "common/mmapReader.h" + +#include +#include +#include +#include + +namespace trt_edgellm +{ +namespace rt +{ +namespace +{ + +rt::Coords coordsFromJson(nlohmann::json const& shape) +{ + std::vector dims; + dims.reserve(shape.size()); + for (auto const& dim : shape) + { + dims.push_back(dim.get()); + } + return rt::Coords(dims); +} + +nvinfer1::DataType dataTypeFromTorchString(std::string const& dtype) +{ + if (dtype == "torch.float16" || dtype == "float16" || dtype == "fp16") + { + return nvinfer1::DataType::kHALF; + } + if (dtype == "torch.float32" || dtype == "float32" || dtype == "fp32") + { + return nvinfer1::DataType::kFLOAT; + } + throw std::runtime_error("MotBackboneRunner: unsupported tensor dtype: " + dtype); +} + +std::string resolveIOTensorName(nvinfer1::ICudaEngine const* engine, std::string const& preferredName, + nvinfer1::TensorIOMode mode, int32_t modeIndex = 0) +{ + for (int32_t i = 0; i < engine->getNbIOTensors(); ++i) + { + auto const* name = engine->getIOTensorName(i); + if (name != nullptr && preferredName == name && engine->getTensorIOMode(name) == mode) + { + return preferredName; + } + } + + int32_t seen = 0; + for (int32_t i = 0; i < engine->getNbIOTensors(); ++i) + { + auto const* name = engine->getIOTensorName(i); + if (name != nullptr && engine->getTensorIOMode(name) == mode) + { + if (seen == modeIndex) + { + LOG_WARNING( + "MotBackboneRunner: tensor name '%s' is not present in the TensorRT engine; using binding '%s'", + preferredName.c_str(), name); + return name; + } + ++seen; + } + } + + throw std::runtime_error("MotBackboneRunner: failed to resolve TensorRT I/O tensor name: " + preferredName); +} + +std::string resolveEnginePath(std::string const& engineDir, nlohmann::json const& configJson) +{ + std::string const engineFile = configJson.value("engine_file", std::string{"mot_backbone.engine"}); + std::filesystem::path const preferred = std::filesystem::path(engineDir) / engineFile; + if (std::filesystem::exists(preferred)) + { + return preferred.string(); + } + + for (char const* candidate : {"mot_backbone.engine", "backbone.engine"}) + { + std::filesystem::path const path = std::filesystem::path(engineDir) / candidate; + if (std::filesystem::exists(path)) + { + return path.string(); + } + } + + return preferred.string(); +} + +int64_t tensorBytes(rt::Tensor const& tensor) +{ + return static_cast(tensor.getShape().volume()) * rt::utils::getTypeSize(tensor.getDataType()); +} + +bool copyTensorToDevice(rt::Tensor& dst, rt::Tensor const& src, cudaStream_t stream) +{ + auto const dstBytes = tensorBytes(dst); + auto const srcBytes = tensorBytes(src); + if (dstBytes == srcBytes) + { + CUDA_CHECK(cudaMemcpyAsync( + dst.rawPointer(), src.rawPointer(), static_cast(dstBytes), cudaMemcpyDeviceToDevice, stream)); + return true; + } + + auto const elements = dst.getShape().volume(); + if (src.getShape().volume() != elements) + { + LOG_ERROR("MotBackboneRunner: copy shape mismatch, dst=%ld src=%ld", elements, src.getShape().volume()); + return false; + } + + auto const srcType = src.getDataType(); + auto const dstType = dst.getDataType(); + if (srcType == nvinfer1::DataType::kHALF && dstType == nvinfer1::DataType::kFLOAT) + { + std::vector hostSrc(static_cast(elements)); + CUDA_CHECK(cudaMemcpyAsync( + hostSrc.data(), src.rawPointer(), static_cast(srcBytes), cudaMemcpyDeviceToHost, stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + std::vector hostDst(static_cast(elements)); + for (int64_t i = 0; i < elements; ++i) + { + hostDst[static_cast(i)] = __half2float(hostSrc[static_cast(i)]); + } + CUDA_CHECK(cudaMemcpyAsync( + dst.rawPointer(), hostDst.data(), static_cast(dstBytes), cudaMemcpyHostToDevice, stream)); + return true; + } + + if (srcType == nvinfer1::DataType::kFLOAT && dstType == nvinfer1::DataType::kHALF) + { + std::vector hostSrc(static_cast(elements)); + CUDA_CHECK(cudaMemcpyAsync( + hostSrc.data(), src.rawPointer(), static_cast(srcBytes), cudaMemcpyDeviceToHost, stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + std::vector hostDst(static_cast(elements)); + for (int64_t i = 0; i < elements; ++i) + { + hostDst[static_cast(i)] = __float2half(hostSrc[static_cast(i)]); + } + CUDA_CHECK(cudaMemcpyAsync( + dst.rawPointer(), hostDst.data(), static_cast(dstBytes), cudaMemcpyHostToDevice, stream)); + return true; + } + + LOG_ERROR("MotBackboneRunner: unsupported dtype conversion src=%d dst=%d", static_cast(srcType), + static_cast(dstType)); + return false; +} + +} // namespace + +MotBackboneRunner::MotBackboneRunner(std::string const& engineDir, cudaStream_t stream) + : mEngineDir(engineDir) +{ + if (!loadConfig()) + { + throw std::runtime_error("MotBackboneRunner: failed to load config from " + engineDir); + } + if (!loadEngine(stream)) + { + throw std::runtime_error("MotBackboneRunner: failed to load TensorRT engine from " + engineDir); + } + if (!validateAndFillConfig()) + { + throw std::runtime_error("MotBackboneRunner: failed to validate config for " + engineDir); + } + if (!allocateBuffers()) + { + throw std::runtime_error("MotBackboneRunner: failed to allocate buffers for " + engineDir); + } + + LOG_INFO("MotBackboneRunner loaded from %s (und_seq=%s, gen_seq=%s, last_hidden=%s)", engineDir.c_str(), + mUndSeqShape.formatString().c_str(), mGenSeqShape.formatString().c_str(), mOutputShape.formatString().c_str()); +} + +bool MotBackboneRunner::loadConfig() +{ + auto const configPath = mEngineDir + "/config.json"; + std::ifstream configFile(configPath); + if (!configFile) + { + LOG_ERROR("MotBackboneRunner: failed to open config file: %s", configPath.c_str()); + return false; + } + + try + { + configFile >> mConfigJson; + } + catch (nlohmann::json::parse_error const& e) + { + LOG_ERROR("MotBackboneRunner: failed to parse config file: %s", e.what()); + return false; + } + return true; +} + +bool MotBackboneRunner::loadEngine(cudaStream_t stream) +{ + auto const enginePath = resolveEnginePath(mEngineDir, mConfigJson); + if (!std::filesystem::exists(enginePath)) + { + LOG_ERROR("MotBackboneRunner: engine not found at %s", enginePath.c_str()); + return false; + } + + file_io::MmapReader engineFileReader(enginePath); + mRuntime.reset(nvinfer1::createInferRuntime(gLogger)); + if (!mRuntime) + { + LOG_ERROR("MotBackboneRunner: failed to create TensorRT runtime for %s", enginePath.c_str()); + return false; + } + + mEngine.reset(mRuntime->deserializeCudaEngine(engineFileReader.getData(), engineFileReader.getSize())); + if (!mEngine) + { + LOG_ERROR("MotBackboneRunner: failed to deserialize TensorRT engine: %s", enginePath.c_str()); + return false; + } + + mContext.reset(mEngine->createExecutionContext(nvinfer1::ExecutionContextAllocationStrategy::kUSER_MANAGED)); + if (!mContext) + { + LOG_ERROR("MotBackboneRunner: failed to create TensorRT execution context for %s", enginePath.c_str()); + return false; + } + + if (!mContext->setOptimizationProfileAsync(0, stream)) + { + LOG_ERROR("MotBackboneRunner: failed to set optimization profile for %s", enginePath.c_str()); + return false; + } + CUDA_CHECK(cudaStreamSynchronize(stream)); + return true; +} + +bool MotBackboneRunner::validateAndFillConfig() +{ + auto const modelType = mConfigJson.value("model_type", std::string{"cosmos_mot_backbone"}); + auto const component = mConfigJson.value("component", std::string{"mot_backbone"}); + if (modelType != "cosmos_mot_backbone" && component != "mot_backbone") + { + LOG_ERROR("MotBackboneRunner: unexpected model_type=%s component=%s", modelType.c_str(), component.c_str()); + return false; + } + + if (mConfigJson.contains("input_names") && mConfigJson.at("input_names").size() >= 6U) + { + auto const& names = mConfigJson.at("input_names"); + mUndSeqInputName = names.at(0).get(); + mGenSeqInputName = names.at(1).get(); + mCosUndInputName = names.at(2).get(); + mSinUndInputName = names.at(3).get(); + mCosGenInputName = names.at(4).get(); + mSinGenInputName = names.at(5).get(); + } + if (mConfigJson.contains("output_names") && !mConfigJson.at("output_names").empty()) + { + mOutputName = mConfigJson.at("output_names").at(0).get(); + } + + if (!mConfigJson.contains("inputs")) + { + LOG_ERROR("MotBackboneRunner: config is missing inputs metadata"); + return false; + } + if (!mConfigJson.contains("outputs") || mConfigJson.at("outputs").empty()) + { + LOG_ERROR("MotBackboneRunner: config is missing output metadata"); + return false; + } + + auto const& inputs = mConfigJson.at("inputs"); + for (auto const* name : {&mUndSeqInputName, &mGenSeqInputName, &mCosUndInputName, &mSinUndInputName, + &mCosGenInputName, &mSinGenInputName}) + { + if (!inputs.contains(*name)) + { + LOG_ERROR("MotBackboneRunner: config is missing input metadata for %s", name->c_str()); + return false; + } + } + + auto const& undMeta = inputs.at(mUndSeqInputName); + auto const& genMeta = inputs.at(mGenSeqInputName); + auto const& cosUndMeta = inputs.at(mCosUndInputName); + auto const& cosGenMeta = inputs.at(mCosGenInputName); + auto const& outputMeta = mConfigJson.at("outputs").at(0); + + mUndSeqShape = coordsFromJson(undMeta.at("shape")); + mGenSeqShape = coordsFromJson(genMeta.at("shape")); + mRotaryUndShape = coordsFromJson(cosUndMeta.at("shape")); + mRotaryGenShape = coordsFromJson(cosGenMeta.at("shape")); + mOutputShape = coordsFromJson(outputMeta.at("shape")); + + mSeqType = dataTypeFromTorchString(undMeta.at("dtype").get()); + mRotaryType = dataTypeFromTorchString(cosUndMeta.at("dtype").get()); + mOutputType = dataTypeFromTorchString(outputMeta.at("dtype").get()); + + if (mUndSeqShape.getNumDims() != 2 || mGenSeqShape.getNumDims() != 2) + { + LOG_ERROR("MotBackboneRunner: expected rank-2 und_seq and gen_seq tensors"); + return false; + } + if (mRotaryUndShape.getNumDims() != 2 || mRotaryGenShape.getNumDims() != 2) + { + LOG_ERROR("MotBackboneRunner: expected rank-2 rotary tensors"); + return false; + } + if (mOutputShape.getNumDims() != 2) + { + LOG_ERROR("MotBackboneRunner: expected rank-2 last_hidden_state output"); + return false; + } + if (mUndSeqShape[0] + mGenSeqShape[0] != mOutputShape[0]) + { + LOG_ERROR("MotBackboneRunner: output seq dim %ld != und_len %ld + gen_len %ld", mOutputShape[0], + mUndSeqShape[0], mGenSeqShape[0]); + return false; + } + + mUndSeqInputName = resolveIOTensorName(mEngine.get(), mUndSeqInputName, nvinfer1::TensorIOMode::kINPUT, 0); + mGenSeqInputName = resolveIOTensorName(mEngine.get(), mGenSeqInputName, nvinfer1::TensorIOMode::kINPUT, 1); + mCosUndInputName = resolveIOTensorName(mEngine.get(), mCosUndInputName, nvinfer1::TensorIOMode::kINPUT, 2); + mSinUndInputName = resolveIOTensorName(mEngine.get(), mSinUndInputName, nvinfer1::TensorIOMode::kINPUT, 3); + mCosGenInputName = resolveIOTensorName(mEngine.get(), mCosGenInputName, nvinfer1::TensorIOMode::kINPUT, 4); + mSinGenInputName = resolveIOTensorName(mEngine.get(), mSinGenInputName, nvinfer1::TensorIOMode::kINPUT, 5); + mOutputName = resolveIOTensorName(mEngine.get(), mOutputName, nvinfer1::TensorIOMode::kOUTPUT, 0); + return true; +} + +int64_t MotBackboneRunner::getRequiredContextMemorySize() const +{ + return mEngine ? mEngine->getDeviceMemorySizeV2() : 0; +} + +bool MotBackboneRunner::allocateBuffers() +{ + mUndSeqTensor = rt::Tensor(mUndSeqShape, rt::DeviceType::kGPU, mSeqType, mUndSeqInputName); + mGenSeqTensor = rt::Tensor(mGenSeqShape, rt::DeviceType::kGPU, mSeqType, mGenSeqInputName); + mCosUndTensor = rt::Tensor(mRotaryUndShape, rt::DeviceType::kGPU, mRotaryType, mCosUndInputName); + mSinUndTensor = rt::Tensor(mRotaryUndShape, rt::DeviceType::kGPU, mRotaryType, mSinUndInputName); + mCosGenTensor = rt::Tensor(mRotaryGenShape, rt::DeviceType::kGPU, mRotaryType, mCosGenInputName); + mSinGenTensor = rt::Tensor(mRotaryGenShape, rt::DeviceType::kGPU, mRotaryType, mSinGenInputName); + mLastHiddenTensor = rt::Tensor(mOutputShape, rt::DeviceType::kGPU, mOutputType, mOutputName); + return bindTensors(); +} + +bool MotBackboneRunner::bindTensors() noexcept +{ + struct InputBinding + { + std::string const& name; + rt::Tensor& tensor; + }; + + InputBinding const bindings[] + = {{mUndSeqInputName, mUndSeqTensor}, {mGenSeqInputName, mGenSeqTensor}, {mCosUndInputName, mCosUndTensor}, + {mSinUndInputName, mSinUndTensor}, {mCosGenInputName, mCosGenTensor}, {mSinGenInputName, mSinGenTensor}}; + + for (auto const& binding : bindings) + { + if (!mContext->setInputShape(binding.name.c_str(), binding.tensor.getShape().getTRTDims())) + { + LOG_ERROR("MotBackboneRunner: failed to set input shape for %s", binding.name.c_str()); + return false; + } + if (!mContext->setTensorAddress(binding.name.c_str(), binding.tensor.rawPointer())) + { + LOG_ERROR("MotBackboneRunner: failed to bind input tensor %s", binding.name.c_str()); + return false; + } + } + + if (!mContext->setTensorAddress(mOutputName.c_str(), mLastHiddenTensor.rawPointer())) + { + LOG_ERROR("MotBackboneRunner: failed to bind output tensor %s", mOutputName.c_str()); + return false; + } + return true; +} + +bool MotBackboneRunner::setContextMemory(rt::Tensor& sharedContextMemory) +{ + if (!mEngine) + { + return true; + } + + int64_t const requiredSize = getRequiredContextMemorySize(); + if (sharedContextMemory.getMemoryCapacity() < requiredSize) + { + LOG_ERROR("MotBackboneRunner: shared context memory (%zu bytes) is smaller than required (%zu bytes)", + static_cast(sharedContextMemory.getMemoryCapacity()), static_cast(requiredSize)); + return false; + } + + mContext->setDeviceMemoryV2(sharedContextMemory.rawPointer(), sharedContextMemory.getMemoryCapacity()); + return true; +} + +bool MotBackboneRunner::copyTensorInput( + rt::Tensor& dst, rt::Tensor const& src, char const* tensorLabel, cudaStream_t stream) +{ + if (src.getShape() != dst.getShape()) + { + LOG_ERROR("MotBackboneRunner: %s shape %s does not match engine input %s", tensorLabel, + src.getShape().formatString().c_str(), dst.getShape().formatString().c_str()); + return false; + } + return copyTensorToDevice(dst, src, stream); +} + +bool MotBackboneRunner::copyUndSeqFrom(rt::Tensor const& src, cudaStream_t stream) +{ + return copyTensorInput(mUndSeqTensor, src, "und_seq", stream); +} + +bool MotBackboneRunner::copyGenSeqFrom(rt::Tensor const& src, cudaStream_t stream) +{ + return copyTensorInput(mGenSeqTensor, src, "gen_seq", stream); +} + +bool MotBackboneRunner::copyRotaryFrom(CosmosTextPhase0 const& phase0, cudaStream_t stream) +{ + rt::Coords const expectedCosUnd({phase0.undLen, mRotaryUndShape[1]}); + rt::Coords const expectedSinUnd({phase0.undLen, mRotaryUndShape[1]}); + rt::Coords const expectedCosGen({phase0.genLen, mRotaryGenShape[1]}); + rt::Coords const expectedSinGen({phase0.genLen, mRotaryGenShape[1]}); + + if (expectedCosUnd != mRotaryUndShape || expectedSinUnd != mRotaryUndShape || expectedCosGen != mRotaryGenShape + || expectedSinGen != mRotaryGenShape) + { + LOG_ERROR( + "MotBackboneRunner: rotary shapes from Phase 0 (und=%s gen=%s) do not match engine bindings (und=%s " + "gen=%s).", + expectedCosUnd.formatString().c_str(), expectedCosGen.formatString().c_str(), + mRotaryUndShape.formatString().c_str(), mRotaryGenShape.formatString().c_str()); + return false; + } + + return copyTensorInput(mCosUndTensor, phase0.cosUnd, "cos_und", stream) + && copyTensorInput(mSinUndTensor, phase0.sinUnd, "sin_und", stream) + && copyTensorInput(mCosGenTensor, phase0.cosGen, "cos_gen", stream) + && copyTensorInput(mSinGenTensor, phase0.sinGen, "sin_gen", stream); +} + +bool MotBackboneRunner::runBackbone(cudaStream_t stream) noexcept +{ + if (!bindTensors()) + { + return false; + } + if (!mContext->enqueueV3(stream)) + { + LOG_ERROR("MotBackboneRunner: enqueueV3 failed."); + return false; + } + return true; +} + +} // namespace rt +} // namespace trt_edgellm diff --git a/cpp/runtime/motBackboneRunner.h b/cpp/runtime/motBackboneRunner.h new file mode 100644 index 00000000..25eb0527 --- /dev/null +++ b/cpp/runtime/motBackboneRunner.h @@ -0,0 +1,191 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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. + */ + +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "common/tensor.h" +#include "runtime/wfmRuntimeUtils.h" + +#include +#include +#include + +#include +#include +#include + +namespace trt_edgellm +{ +namespace rt +{ + +//! Cosmos MoT backbone engine: und_seq + gen_seq + rotary -> last_hidden_state. +class MotBackboneRunner +{ +public: + //! \p engineDir is the ``mot_backbone/`` component directory (contains config.json + engine). + explicit MotBackboneRunner(std::string const& engineDir, cudaStream_t stream); + + ~MotBackboneRunner() noexcept = default; + + int64_t getRequiredContextMemorySize() const; + bool setContextMemory(rt::Tensor& sharedContextMemory); + + std::string const& getOutputName() const noexcept + { + return mOutputName; + } + + rt::Coords const& getUndSeqShape() const noexcept + { + return mUndSeqShape; + } + + rt::Coords const& getGenSeqShape() const noexcept + { + return mGenSeqShape; + } + + rt::Coords const& getLastHiddenShape() const noexcept + { + return mOutputShape; + } + + rt::Tensor& getUndSeq() noexcept + { + return mUndSeqTensor; + } + + rt::Tensor const& getUndSeq() const noexcept + { + return mUndSeqTensor; + } + + rt::Tensor& getGenSeq() noexcept + { + return mGenSeqTensor; + } + + rt::Tensor const& getGenSeq() const noexcept + { + return mGenSeqTensor; + } + + rt::Tensor& getCosUnd() noexcept + { + return mCosUndTensor; + } + + rt::Tensor const& getCosUnd() const noexcept + { + return mCosUndTensor; + } + + rt::Tensor& getSinUnd() noexcept + { + return mSinUndTensor; + } + + rt::Tensor const& getSinUnd() const noexcept + { + return mSinUndTensor; + } + + rt::Tensor& getCosGen() noexcept + { + return mCosGenTensor; + } + + rt::Tensor const& getCosGen() const noexcept + { + return mCosGenTensor; + } + + rt::Tensor& getSinGen() noexcept + { + return mSinGenTensor; + } + + rt::Tensor const& getSinGen() const noexcept + { + return mSinGenTensor; + } + + rt::Tensor& getLastHiddenState() noexcept + { + return mLastHiddenTensor; + } + + rt::Tensor const& getLastHiddenState() const noexcept + { + return mLastHiddenTensor; + } + + bool copyUndSeqFrom(rt::Tensor const& src, cudaStream_t stream); + bool copyGenSeqFrom(rt::Tensor const& src, cudaStream_t stream); + bool copyRotaryFrom(CosmosTextPhase0 const& phase0, cudaStream_t stream); + bool runBackbone(cudaStream_t stream) noexcept; + +private: + bool loadConfig(); + bool loadEngine(cudaStream_t stream); + bool validateAndFillConfig(); + bool allocateBuffers(); + bool bindTensors() noexcept; + bool copyTensorInput(rt::Tensor& dst, rt::Tensor const& src, char const* tensorLabel, cudaStream_t stream); + + std::string mEngineDir; + nlohmann::json mConfigJson; + + std::unique_ptr mRuntime; + std::unique_ptr mEngine; + std::unique_ptr mContext; + + std::string mUndSeqInputName{"und_seq"}; + std::string mGenSeqInputName{"gen_seq"}; + std::string mCosUndInputName{"cos_und"}; + std::string mSinUndInputName{"sin_und"}; + std::string mCosGenInputName{"cos_gen"}; + std::string mSinGenInputName{"sin_gen"}; + std::string mOutputName{"last_hidden_state"}; + + rt::Coords mUndSeqShape; + rt::Coords mGenSeqShape; + rt::Coords mRotaryUndShape; + rt::Coords mRotaryGenShape; + rt::Coords mOutputShape; + + nvinfer1::DataType mSeqType{nvinfer1::DataType::kHALF}; + nvinfer1::DataType mRotaryType{nvinfer1::DataType::kHALF}; + nvinfer1::DataType mOutputType{nvinfer1::DataType::kHALF}; + + rt::Tensor mUndSeqTensor; + rt::Tensor mGenSeqTensor; + rt::Tensor mCosUndTensor; + rt::Tensor mSinUndTensor; + rt::Tensor mCosGenTensor; + rt::Tensor mSinGenTensor; + rt::Tensor mLastHiddenTensor; +}; + +} // namespace rt +} // namespace trt_edgellm diff --git a/cpp/runtime/vaeDecodeRunner.cpp b/cpp/runtime/vaeDecodeRunner.cpp new file mode 100644 index 00000000..bc13bf00 --- /dev/null +++ b/cpp/runtime/vaeDecodeRunner.cpp @@ -0,0 +1,416 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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. + */ + +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/vaeDecodeRunner.h" + +#include "common/cudaUtils.h" +#include "common/logger.h" +#include "common/mmapReader.h" + +#include +#include +#include +#include + +namespace trt_edgellm +{ +namespace rt +{ +namespace +{ + +rt::Coords coordsFromJson(nlohmann::json const& shape) +{ + std::vector dims; + dims.reserve(shape.size()); + for (auto const& dim : shape) + { + dims.push_back(dim.get()); + } + return rt::Coords(dims); +} + +nvinfer1::DataType dataTypeFromTorchString(std::string const& dtype) +{ + if (dtype == "torch.float16" || dtype == "float16" || dtype == "fp16") + { + return nvinfer1::DataType::kHALF; + } + if (dtype == "torch.float32" || dtype == "float32" || dtype == "fp32") + { + return nvinfer1::DataType::kFLOAT; + } + throw std::runtime_error("VaeDecodeRunner: unsupported tensor dtype: " + dtype); +} + +std::string resolveIOTensorName(nvinfer1::ICudaEngine const* engine, std::string const& preferredName, + nvinfer1::TensorIOMode mode, int32_t modeIndex = 0) +{ + for (int32_t i = 0; i < engine->getNbIOTensors(); ++i) + { + auto const* name = engine->getIOTensorName(i); + if (name != nullptr && preferredName == name && engine->getTensorIOMode(name) == mode) + { + return preferredName; + } + } + + int32_t seen = 0; + for (int32_t i = 0; i < engine->getNbIOTensors(); ++i) + { + auto const* name = engine->getIOTensorName(i); + if (name != nullptr && engine->getTensorIOMode(name) == mode) + { + if (seen == modeIndex) + { + LOG_WARNING( + "VaeDecodeRunner: tensor name '%s' is not present in the TensorRT engine; using binding '%s'", + preferredName.c_str(), name); + return name; + } + ++seen; + } + } + + throw std::runtime_error("VaeDecodeRunner: failed to resolve TensorRT I/O tensor name: " + preferredName); +} + +std::string resolveEnginePath(std::string const& engineDir, nlohmann::json const& configJson) +{ + std::string const engineFile = configJson.value("engine_file", std::string{"visual_decode.engine"}); + std::filesystem::path const preferred = std::filesystem::path(engineDir) / engineFile; + if (std::filesystem::exists(preferred)) + { + return preferred.string(); + } + + for (char const* candidate : {"visual_decode.engine", "visual.engine"}) + { + std::filesystem::path const path = std::filesystem::path(engineDir) / candidate; + if (std::filesystem::exists(path)) + { + return path.string(); + } + } + + return preferred.string(); +} + +int64_t tensorBytes(rt::Tensor const& tensor) +{ + return static_cast(tensor.getShape().volume()) * rt::utils::getTypeSize(tensor.getDataType()); +} + +bool copyTensorToDevice(rt::Tensor& dst, rt::Tensor const& src, cudaStream_t stream) +{ + auto const dstBytes = tensorBytes(dst); + auto const srcBytes = tensorBytes(src); + if (dstBytes == srcBytes) + { + CUDA_CHECK(cudaMemcpyAsync( + dst.rawPointer(), src.rawPointer(), static_cast(dstBytes), cudaMemcpyDeviceToDevice, stream)); + return true; + } + + auto const elements = dst.getShape().volume(); + if (src.getShape().volume() != elements) + { + LOG_ERROR("VaeDecodeRunner: copy shape mismatch, dst=%ld src=%ld", elements, src.getShape().volume()); + return false; + } + + auto const srcType = src.getDataType(); + auto const dstType = dst.getDataType(); + if (srcType == nvinfer1::DataType::kHALF && dstType == nvinfer1::DataType::kFLOAT) + { + std::vector hostSrc(static_cast(elements)); + CUDA_CHECK(cudaMemcpyAsync( + hostSrc.data(), src.rawPointer(), static_cast(srcBytes), cudaMemcpyDeviceToHost, stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + std::vector hostDst(static_cast(elements)); + for (int64_t i = 0; i < elements; ++i) + { + hostDst[static_cast(i)] = __half2float(hostSrc[static_cast(i)]); + } + CUDA_CHECK(cudaMemcpyAsync( + dst.rawPointer(), hostDst.data(), static_cast(dstBytes), cudaMemcpyHostToDevice, stream)); + return true; + } + + if (srcType == nvinfer1::DataType::kFLOAT && dstType == nvinfer1::DataType::kHALF) + { + std::vector hostSrc(static_cast(elements)); + CUDA_CHECK(cudaMemcpyAsync( + hostSrc.data(), src.rawPointer(), static_cast(srcBytes), cudaMemcpyDeviceToHost, stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + std::vector hostDst(static_cast(elements)); + for (int64_t i = 0; i < elements; ++i) + { + hostDst[static_cast(i)] = __float2half(hostSrc[static_cast(i)]); + } + CUDA_CHECK(cudaMemcpyAsync( + dst.rawPointer(), hostDst.data(), static_cast(dstBytes), cudaMemcpyHostToDevice, stream)); + return true; + } + + LOG_ERROR("VaeDecodeRunner: unsupported dtype conversion src=%d dst=%d", static_cast(srcType), + static_cast(dstType)); + return false; +} + +} // namespace + +VaeDecodeRunner::VaeDecodeRunner(std::string const& engineDir, cudaStream_t stream) + : mEngineDir(engineDir) +{ + if (!loadConfig()) + { + throw std::runtime_error("VaeDecodeRunner: failed to load config from " + engineDir); + } + if (!loadEngine(stream)) + { + throw std::runtime_error("VaeDecodeRunner: failed to load TensorRT engine from " + engineDir); + } + if (!validateAndFillConfig()) + { + throw std::runtime_error("VaeDecodeRunner: failed to validate config for " + engineDir); + } + if (!allocateBuffers()) + { + throw std::runtime_error("VaeDecodeRunner: failed to allocate buffers for " + engineDir); + } + + LOG_INFO("VaeDecodeRunner loaded from %s (%s -> %s, latents=%s, pixels=%s)", engineDir.c_str(), mInputName.c_str(), + mOutputName.c_str(), mInputShape.formatString().c_str(), mOutputShape.formatString().c_str()); +} + +bool VaeDecodeRunner::loadConfig() +{ + auto const configPath = mEngineDir + "/config.json"; + std::ifstream configFile(configPath); + if (!configFile) + { + LOG_ERROR("VaeDecodeRunner: failed to open config file: %s", configPath.c_str()); + return false; + } + + try + { + configFile >> mConfigJson; + } + catch (nlohmann::json::parse_error const& e) + { + LOG_ERROR("VaeDecodeRunner: failed to parse config file: %s", e.what()); + return false; + } + return true; +} + +bool VaeDecodeRunner::loadEngine(cudaStream_t stream) +{ + auto const enginePath = resolveEnginePath(mEngineDir, mConfigJson); + if (!std::filesystem::exists(enginePath)) + { + LOG_ERROR("VaeDecodeRunner: engine not found at %s", enginePath.c_str()); + return false; + } + + file_io::MmapReader engineFileReader(enginePath); + mRuntime.reset(nvinfer1::createInferRuntime(gLogger)); + if (!mRuntime) + { + LOG_ERROR("VaeDecodeRunner: failed to create TensorRT runtime for %s", enginePath.c_str()); + return false; + } + + mEngine.reset(mRuntime->deserializeCudaEngine(engineFileReader.getData(), engineFileReader.getSize())); + if (!mEngine) + { + LOG_ERROR("VaeDecodeRunner: failed to deserialize TensorRT engine: %s", enginePath.c_str()); + return false; + } + + mContext.reset(mEngine->createExecutionContext(nvinfer1::ExecutionContextAllocationStrategy::kUSER_MANAGED)); + if (!mContext) + { + LOG_ERROR("VaeDecodeRunner: failed to create TensorRT execution context for %s", enginePath.c_str()); + return false; + } + + if (!mContext->setOptimizationProfileAsync(0, stream)) + { + LOG_ERROR("VaeDecodeRunner: failed to set optimization profile for %s", enginePath.c_str()); + return false; + } + CUDA_CHECK(cudaStreamSynchronize(stream)); + return true; +} + +bool VaeDecodeRunner::validateAndFillConfig() +{ + auto const modelType = mConfigJson.value("model_type", std::string{"cosmos_vae_decode"}); + auto const component = mConfigJson.value("component", std::string{"visual_decode"}); + if (modelType != "cosmos_vae_decode" && component != "visual_decode") + { + LOG_ERROR("VaeDecodeRunner: unexpected model_type=%s component=%s", modelType.c_str(), component.c_str()); + return false; + } + + if (mConfigJson.contains("input_names") && !mConfigJson.at("input_names").empty()) + { + mInputName = mConfigJson.at("input_names").at(0).get(); + } + if (mConfigJson.contains("output_names") && !mConfigJson.at("output_names").empty()) + { + mOutputName = mConfigJson.at("output_names").at(0).get(); + } + + if (!mConfigJson.contains("inputs") || !mConfigJson.at("inputs").contains(mInputName)) + { + LOG_ERROR("VaeDecodeRunner: config is missing input metadata for %s", mInputName.c_str()); + return false; + } + if (!mConfigJson.contains("outputs") || mConfigJson.at("outputs").empty()) + { + LOG_ERROR("VaeDecodeRunner: config is missing output metadata"); + return false; + } + + auto const& inputMeta = mConfigJson.at("inputs").at(mInputName); + auto const& outputMeta = mConfigJson.at("outputs").at(0); + + mInputShape = coordsFromJson(inputMeta.at("shape")); + mOutputShape = coordsFromJson(outputMeta.at("shape")); + mInputType = dataTypeFromTorchString(inputMeta.at("dtype").get()); + mOutputType = dataTypeFromTorchString(outputMeta.at("dtype").get()); + + if (mInputShape.getNumDims() != 5 || mOutputShape.getNumDims() != 5) + { + LOG_ERROR("VaeDecodeRunner: expected rank-5 latent/pixel tensors (got input=%d output=%d dims)", + mInputShape.getNumDims(), mOutputShape.getNumDims()); + return false; + } + + mInputName = resolveIOTensorName(mEngine.get(), mInputName, nvinfer1::TensorIOMode::kINPUT, 0); + mOutputName = resolveIOTensorName(mEngine.get(), mOutputName, nvinfer1::TensorIOMode::kOUTPUT, 0); + return true; +} + +int64_t VaeDecodeRunner::getRequiredContextMemorySize() const +{ + return mEngine ? mEngine->getDeviceMemorySizeV2() : 0; +} + +bool VaeDecodeRunner::allocateBuffers() +{ + mLatentsTensor = rt::Tensor(mInputShape, rt::DeviceType::kGPU, mInputType, mInputName); + mPixelsTensor = rt::Tensor(mOutputShape, rt::DeviceType::kGPU, mOutputType, mOutputName); + return bindTensors(); +} + +bool VaeDecodeRunner::bindTensors() noexcept +{ + if (!mContext->setInputShape(mInputName.c_str(), mLatentsTensor.getShape().getTRTDims())) + { + LOG_ERROR("VaeDecodeRunner: failed to set input shape for %s", mInputName.c_str()); + return false; + } + if (!mContext->setTensorAddress(mInputName.c_str(), mLatentsTensor.rawPointer())) + { + LOG_ERROR("VaeDecodeRunner: failed to bind input tensor %s", mInputName.c_str()); + return false; + } + if (!mContext->setTensorAddress(mOutputName.c_str(), mPixelsTensor.rawPointer())) + { + LOG_ERROR("VaeDecodeRunner: failed to bind output tensor %s", mOutputName.c_str()); + return false; + } + return true; +} + +bool VaeDecodeRunner::setContextMemory(rt::Tensor& sharedContextMemory) +{ + if (!mEngine) + { + return true; + } + + int64_t const requiredSize = getRequiredContextMemorySize(); + if (sharedContextMemory.getMemoryCapacity() < requiredSize) + { + LOG_ERROR("VaeDecodeRunner: shared context memory (%zu bytes) is smaller than required (%zu bytes)", + static_cast(sharedContextMemory.getMemoryCapacity()), static_cast(requiredSize)); + return false; + } + + mContext->setDeviceMemoryV2(sharedContextMemory.rawPointer(), sharedContextMemory.getMemoryCapacity()); + return true; +} + +bool VaeDecodeRunner::copyLatentsFrom(rt::Tensor const& src, cudaStream_t stream) +{ + if (src.getShape() != mLatentsTensor.getShape()) + { + LOG_ERROR("VaeDecodeRunner: latent shape %s does not match engine input %s", + src.getShape().formatString().c_str(), mLatentsTensor.getShape().formatString().c_str()); + return false; + } + return copyTensorToDevice(mLatentsTensor, src, stream); +} + +bool VaeDecodeRunner::copyPixelsTo(rt::Tensor& dst, cudaStream_t stream) const +{ + if (dst.getShape() != mPixelsTensor.getShape()) + { + LOG_ERROR("VaeDecodeRunner: pixel shape %s does not match engine output %s", + dst.getShape().formatString().c_str(), mPixelsTensor.getShape().formatString().c_str()); + return false; + } + return copyTensorToDevice(dst, mPixelsTensor, stream); +} + +bool VaeDecodeRunner::copyPixelsTo(VideoBuffer& video, cudaStream_t stream) const +{ + if (!video.buffer) + { + LOG_ERROR("VaeDecodeRunner: VideoBuffer is missing a tensor."); + return false; + } + return copyPixelsTo(*video.buffer, stream); +} + +bool VaeDecodeRunner::decode(cudaStream_t stream) noexcept +{ + if (!bindTensors()) + { + return false; + } + if (!mContext->enqueueV3(stream)) + { + LOG_ERROR("VaeDecodeRunner: enqueueV3 failed."); + return false; + } + return true; +} + +} // namespace rt +} // namespace trt_edgellm diff --git a/cpp/runtime/vaeDecodeRunner.h b/cpp/runtime/vaeDecodeRunner.h new file mode 100644 index 00000000..dc2f0567 --- /dev/null +++ b/cpp/runtime/vaeDecodeRunner.h @@ -0,0 +1,124 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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. + */ + +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "common/tensor.h" +#include "runtime/wfmRuntimeUtils.h" + +#include +#include +#include + +#include +#include +#include + +namespace trt_edgellm +{ +namespace rt +{ + +//! Wan VAE decode engine: latents [B,C,T',H',W'] -> pixels [B,3,T,H,W]. +class VaeDecodeRunner +{ +public: + //! \p engineDir is the ``visual_decode/`` component directory (contains config.json + engine). + explicit VaeDecodeRunner(std::string const& engineDir, cudaStream_t stream); + + ~VaeDecodeRunner() noexcept = default; + + int64_t getRequiredContextMemorySize() const; + bool setContextMemory(rt::Tensor& sharedContextMemory); + + std::string const& getInputName() const noexcept + { + return mInputName; + } + + std::string const& getOutputName() const noexcept + { + return mOutputName; + } + + rt::Coords const& getLatentShape() const noexcept + { + return mInputShape; + } + + rt::Coords const& getPixelShape() const noexcept + { + return mOutputShape; + } + + rt::Tensor& getLatents() noexcept + { + return mLatentsTensor; + } + + rt::Tensor const& getLatents() const noexcept + { + return mLatentsTensor; + } + + rt::Tensor& getPixels() noexcept + { + return mPixelsTensor; + } + + rt::Tensor const& getPixels() const noexcept + { + return mPixelsTensor; + } + + bool copyLatentsFrom(rt::Tensor const& src, cudaStream_t stream); + bool copyPixelsTo(rt::Tensor& dst, cudaStream_t stream) const; + bool copyPixelsTo(VideoBuffer& video, cudaStream_t stream) const; + bool decode(cudaStream_t stream) noexcept; + +private: + bool loadConfig(); + bool loadEngine(cudaStream_t stream); + bool validateAndFillConfig(); + bool allocateBuffers(); + bool bindTensors() noexcept; + + std::string mEngineDir; + nlohmann::json mConfigJson; + + std::unique_ptr mRuntime; + std::unique_ptr mEngine; + std::unique_ptr mContext; + + std::string mInputName{"latents"}; + std::string mOutputName{"pixels"}; + rt::Coords mInputShape; + rt::Coords mOutputShape; + nvinfer1::DataType mInputType{nvinfer1::DataType::kHALF}; + nvinfer1::DataType mOutputType{nvinfer1::DataType::kHALF}; + + rt::Tensor mLatentsTensor; + rt::Tensor mPixelsTensor; +}; + +} // namespace rt +} // namespace trt_edgellm diff --git a/cpp/runtime/vaeEncodeRunner.cpp b/cpp/runtime/vaeEncodeRunner.cpp new file mode 100644 index 00000000..1a6d1256 --- /dev/null +++ b/cpp/runtime/vaeEncodeRunner.cpp @@ -0,0 +1,406 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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. + */ + +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/vaeEncodeRunner.h" + +#include "common/checkMacros.h" +#include "common/cudaUtils.h" +#include "common/logger.h" +#include "common/mmapReader.h" + +#include +#include +#include +#include + +namespace trt_edgellm +{ +namespace rt +{ +namespace +{ + +rt::Coords coordsFromJson(nlohmann::json const& shape) +{ + std::vector dims; + dims.reserve(shape.size()); + for (auto const& dim : shape) + { + dims.push_back(dim.get()); + } + return rt::Coords(dims); +} + +nvinfer1::DataType dataTypeFromTorchString(std::string const& dtype) +{ + if (dtype == "torch.float16" || dtype == "float16" || dtype == "fp16") + { + return nvinfer1::DataType::kHALF; + } + if (dtype == "torch.float32" || dtype == "float32" || dtype == "fp32") + { + return nvinfer1::DataType::kFLOAT; + } + throw std::runtime_error("VaeEncodeRunner: unsupported tensor dtype: " + dtype); +} + +std::string resolveIOTensorName(nvinfer1::ICudaEngine const* engine, std::string const& preferredName, + nvinfer1::TensorIOMode mode, int32_t modeIndex = 0) +{ + for (int32_t i = 0; i < engine->getNbIOTensors(); ++i) + { + auto const* name = engine->getIOTensorName(i); + if (name != nullptr && preferredName == name && engine->getTensorIOMode(name) == mode) + { + return preferredName; + } + } + + int32_t seen = 0; + for (int32_t i = 0; i < engine->getNbIOTensors(); ++i) + { + auto const* name = engine->getIOTensorName(i); + if (name != nullptr && engine->getTensorIOMode(name) == mode) + { + if (seen == modeIndex) + { + LOG_WARNING( + "VaeEncodeRunner: tensor name '%s' is not present in the TensorRT engine; using binding '%s'", + preferredName.c_str(), name); + return name; + } + ++seen; + } + } + + throw std::runtime_error("VaeEncodeRunner: failed to resolve TensorRT I/O tensor name: " + preferredName); +} + +std::string resolveEnginePath(std::string const& engineDir, nlohmann::json const& configJson) +{ + std::string const engineFile = configJson.value("engine_file", std::string{"visual_encode.engine"}); + std::filesystem::path const preferred = std::filesystem::path(engineDir) / engineFile; + if (std::filesystem::exists(preferred)) + { + return preferred.string(); + } + + for (char const* candidate : {"visual_encode.engine", "visual.engine"}) + { + std::filesystem::path const path = std::filesystem::path(engineDir) / candidate; + if (std::filesystem::exists(path)) + { + return path.string(); + } + } + + return preferred.string(); +} + +int64_t tensorBytes(rt::Tensor const& tensor) +{ + return static_cast(tensor.getShape().volume()) * rt::utils::getTypeSize(tensor.getDataType()); +} + +bool copyTensorToDevice(rt::Tensor& dst, rt::Tensor const& src, cudaStream_t stream) +{ + auto const dstBytes = tensorBytes(dst); + auto const srcBytes = tensorBytes(src); + if (dstBytes == srcBytes) + { + CUDA_CHECK(cudaMemcpyAsync( + dst.rawPointer(), src.rawPointer(), static_cast(dstBytes), cudaMemcpyDeviceToDevice, stream)); + return true; + } + + auto const elements = dst.getShape().volume(); + if (src.getShape().volume() != elements) + { + LOG_ERROR("VaeEncodeRunner: copy shape mismatch, dst=%ld src=%ld", elements, src.getShape().volume()); + return false; + } + + auto const srcType = src.getDataType(); + auto const dstType = dst.getDataType(); + if (srcType == nvinfer1::DataType::kHALF && dstType == nvinfer1::DataType::kFLOAT) + { + std::vector hostSrc(static_cast(elements)); + CUDA_CHECK(cudaMemcpyAsync( + hostSrc.data(), src.rawPointer(), static_cast(srcBytes), cudaMemcpyDeviceToHost, stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + std::vector hostDst(static_cast(elements)); + for (int64_t i = 0; i < elements; ++i) + { + hostDst[static_cast(i)] = __half2float(hostSrc[static_cast(i)]); + } + CUDA_CHECK(cudaMemcpyAsync( + dst.rawPointer(), hostDst.data(), static_cast(dstBytes), cudaMemcpyHostToDevice, stream)); + return true; + } + + if (srcType == nvinfer1::DataType::kFLOAT && dstType == nvinfer1::DataType::kHALF) + { + std::vector hostSrc(static_cast(elements)); + CUDA_CHECK(cudaMemcpyAsync( + hostSrc.data(), src.rawPointer(), static_cast(srcBytes), cudaMemcpyDeviceToHost, stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + std::vector hostDst(static_cast(elements)); + for (int64_t i = 0; i < elements; ++i) + { + hostDst[static_cast(i)] = __float2half(hostSrc[static_cast(i)]); + } + CUDA_CHECK(cudaMemcpyAsync( + dst.rawPointer(), hostDst.data(), static_cast(dstBytes), cudaMemcpyHostToDevice, stream)); + return true; + } + + LOG_ERROR("VaeEncodeRunner: unsupported dtype conversion src=%d dst=%d", static_cast(srcType), + static_cast(dstType)); + return false; +} + +} // namespace + +VaeEncodeRunner::VaeEncodeRunner(std::string const& engineDir, cudaStream_t stream) + : mEngineDir(engineDir) +{ + if (!loadConfig()) + { + throw std::runtime_error("VaeEncodeRunner: failed to load config from " + engineDir); + } + if (!loadEngine(stream)) + { + throw std::runtime_error("VaeEncodeRunner: failed to load TensorRT engine from " + engineDir); + } + if (!validateAndFillConfig()) + { + throw std::runtime_error("VaeEncodeRunner: failed to validate config for " + engineDir); + } + if (!allocateBuffers()) + { + throw std::runtime_error("VaeEncodeRunner: failed to allocate buffers for " + engineDir); + } + + LOG_INFO("VaeEncodeRunner loaded from %s (%s -> %s, pixels=%s, latents=%s)", engineDir.c_str(), mInputName.c_str(), + mOutputName.c_str(), mInputShape.formatString().c_str(), mOutputShape.formatString().c_str()); +} + +bool VaeEncodeRunner::loadConfig() +{ + auto const configPath = mEngineDir + "/config.json"; + std::ifstream configFile(configPath); + if (!configFile) + { + LOG_ERROR("VaeEncodeRunner: failed to open config file: %s", configPath.c_str()); + return false; + } + + try + { + configFile >> mConfigJson; + } + catch (nlohmann::json::parse_error const& e) + { + LOG_ERROR("VaeEncodeRunner: failed to parse config file: %s", e.what()); + return false; + } + return true; +} + +bool VaeEncodeRunner::loadEngine(cudaStream_t stream) +{ + auto const enginePath = resolveEnginePath(mEngineDir, mConfigJson); + if (!std::filesystem::exists(enginePath)) + { + LOG_ERROR("VaeEncodeRunner: engine not found at %s", enginePath.c_str()); + return false; + } + + file_io::MmapReader engineFileReader(enginePath); + mRuntime.reset(nvinfer1::createInferRuntime(gLogger)); + if (!mRuntime) + { + LOG_ERROR("VaeEncodeRunner: failed to create TensorRT runtime for %s", enginePath.c_str()); + return false; + } + + mEngine.reset(mRuntime->deserializeCudaEngine(engineFileReader.getData(), engineFileReader.getSize())); + if (!mEngine) + { + LOG_ERROR("VaeEncodeRunner: failed to deserialize TensorRT engine: %s", enginePath.c_str()); + return false; + } + + mContext.reset(mEngine->createExecutionContext(nvinfer1::ExecutionContextAllocationStrategy::kUSER_MANAGED)); + if (!mContext) + { + LOG_ERROR("VaeEncodeRunner: failed to create TensorRT execution context for %s", enginePath.c_str()); + return false; + } + + if (!mContext->setOptimizationProfileAsync(0, stream)) + { + LOG_ERROR("VaeEncodeRunner: failed to set optimization profile for %s", enginePath.c_str()); + return false; + } + CUDA_CHECK(cudaStreamSynchronize(stream)); + return true; +} + +bool VaeEncodeRunner::validateAndFillConfig() +{ + auto const modelType = mConfigJson.value("model_type", std::string{"cosmos_vae_encode"}); + auto const component = mConfigJson.value("component", std::string{"visual_encode"}); + if (modelType != "cosmos_vae_encode" && component != "visual_encode") + { + LOG_ERROR("VaeEncodeRunner: unexpected model_type=%s component=%s", modelType.c_str(), component.c_str()); + return false; + } + + if (mConfigJson.contains("input_names") && !mConfigJson.at("input_names").empty()) + { + mInputName = mConfigJson.at("input_names").at(0).get(); + } + if (mConfigJson.contains("output_names") && !mConfigJson.at("output_names").empty()) + { + mOutputName = mConfigJson.at("output_names").at(0).get(); + } + + if (!mConfigJson.contains("inputs") || !mConfigJson.at("inputs").contains(mInputName)) + { + LOG_ERROR("VaeEncodeRunner: config is missing input metadata for %s", mInputName.c_str()); + return false; + } + if (!mConfigJson.contains("outputs") || mConfigJson.at("outputs").empty()) + { + LOG_ERROR("VaeEncodeRunner: config is missing output metadata"); + return false; + } + + auto const& inputMeta = mConfigJson.at("inputs").at(mInputName); + auto const& outputMeta = mConfigJson.at("outputs").at(0); + + mInputShape = coordsFromJson(inputMeta.at("shape")); + mOutputShape = coordsFromJson(outputMeta.at("shape")); + mInputType = dataTypeFromTorchString(inputMeta.at("dtype").get()); + mOutputType = dataTypeFromTorchString(outputMeta.at("dtype").get()); + + if (mInputShape.getNumDims() != 5 || mOutputShape.getNumDims() != 5) + { + LOG_ERROR("VaeEncodeRunner: expected rank-5 pixel/latent tensors (got input=%d output=%d dims)", + mInputShape.getNumDims(), mOutputShape.getNumDims()); + return false; + } + + mInputName = resolveIOTensorName(mEngine.get(), mInputName, nvinfer1::TensorIOMode::kINPUT, 0); + mOutputName = resolveIOTensorName(mEngine.get(), mOutputName, nvinfer1::TensorIOMode::kOUTPUT, 0); + return true; +} + +int64_t VaeEncodeRunner::getRequiredContextMemorySize() const +{ + return mEngine ? mEngine->getDeviceMemorySizeV2() : 0; +} + +bool VaeEncodeRunner::allocateBuffers() +{ + mPixelsTensor = rt::Tensor(mInputShape, rt::DeviceType::kGPU, mInputType, mInputName); + mLatentsTensor = rt::Tensor(mOutputShape, rt::DeviceType::kGPU, mOutputType, mOutputName); + return bindTensors(); +} + +bool VaeEncodeRunner::bindTensors() noexcept +{ + if (!mContext->setInputShape(mInputName.c_str(), mPixelsTensor.getShape().getTRTDims())) + { + LOG_ERROR("VaeEncodeRunner: failed to set input shape for %s", mInputName.c_str()); + return false; + } + if (!mContext->setTensorAddress(mInputName.c_str(), mPixelsTensor.rawPointer())) + { + LOG_ERROR("VaeEncodeRunner: failed to bind input tensor %s", mInputName.c_str()); + return false; + } + if (!mContext->setTensorAddress(mOutputName.c_str(), mLatentsTensor.rawPointer())) + { + LOG_ERROR("VaeEncodeRunner: failed to bind output tensor %s", mOutputName.c_str()); + return false; + } + return true; +} + +bool VaeEncodeRunner::setContextMemory(rt::Tensor& sharedContextMemory) +{ + if (!mEngine) + { + return true; + } + + int64_t const requiredSize = getRequiredContextMemorySize(); + if (sharedContextMemory.getMemoryCapacity() < requiredSize) + { + LOG_ERROR("VaeEncodeRunner: shared context memory (%zu bytes) is smaller than required (%zu bytes)", + static_cast(sharedContextMemory.getMemoryCapacity()), static_cast(requiredSize)); + return false; + } + + mContext->setDeviceMemoryV2(sharedContextMemory.rawPointer(), sharedContextMemory.getMemoryCapacity()); + return true; +} + +bool VaeEncodeRunner::copyPixelsFrom(rt::Tensor const& src, cudaStream_t stream) +{ + if (src.getShape() != mPixelsTensor.getShape()) + { + LOG_ERROR("VaeEncodeRunner: pixel shape %s does not match engine input %s", + src.getShape().formatString().c_str(), mPixelsTensor.getShape().formatString().c_str()); + return false; + } + return copyTensorToDevice(mPixelsTensor, src, stream); +} + +bool VaeEncodeRunner::copyPixelsFrom(VideoBuffer const& video, cudaStream_t stream) +{ + if (!video.buffer) + { + LOG_ERROR("VaeEncodeRunner: VideoBuffer is missing a tensor."); + return false; + } + return copyPixelsFrom(*video.buffer, stream); +} + +bool VaeEncodeRunner::encode(cudaStream_t stream) noexcept +{ + if (!bindTensors()) + { + return false; + } + if (!mContext->enqueueV3(stream)) + { + LOG_ERROR("VaeEncodeRunner: enqueueV3 failed."); + return false; + } + return true; +} + +} // namespace rt +} // namespace trt_edgellm diff --git a/cpp/runtime/vaeEncodeRunner.h b/cpp/runtime/vaeEncodeRunner.h new file mode 100644 index 00000000..bcea227e --- /dev/null +++ b/cpp/runtime/vaeEncodeRunner.h @@ -0,0 +1,123 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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. + */ + +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "common/tensor.h" +#include "runtime/wfmRuntimeUtils.h" + +#include +#include +#include + +#include +#include +#include + +namespace trt_edgellm +{ +namespace rt +{ + +//! Wan VAE encode engine: pixels [B,3,T,H,W] -> latents [B,C,T',H',W']. +class VaeEncodeRunner +{ +public: + //! \p engineDir is the ``visual_encode/`` component directory (contains config.json + engine). + explicit VaeEncodeRunner(std::string const& engineDir, cudaStream_t stream); + + ~VaeEncodeRunner() noexcept = default; + + int64_t getRequiredContextMemorySize() const; + bool setContextMemory(rt::Tensor& sharedContextMemory); + + std::string const& getInputName() const noexcept + { + return mInputName; + } + + std::string const& getOutputName() const noexcept + { + return mOutputName; + } + + rt::Coords const& getPixelShape() const noexcept + { + return mInputShape; + } + + rt::Coords const& getLatentShape() const noexcept + { + return mOutputShape; + } + + rt::Tensor& getPixels() noexcept + { + return mPixelsTensor; + } + + rt::Tensor const& getPixels() const noexcept + { + return mPixelsTensor; + } + + rt::Tensor& getLatents() noexcept + { + return mLatentsTensor; + } + + rt::Tensor const& getLatents() const noexcept + { + return mLatentsTensor; + } + + bool copyPixelsFrom(rt::Tensor const& src, cudaStream_t stream); + bool copyPixelsFrom(VideoBuffer const& video, cudaStream_t stream); + bool encode(cudaStream_t stream) noexcept; + +private: + bool loadConfig(); + bool loadEngine(cudaStream_t stream); + bool validateAndFillConfig(); + bool allocateBuffers(); + bool bindTensors() noexcept; + + std::string mEngineDir; + nlohmann::json mConfigJson; + + std::unique_ptr mRuntime; + std::unique_ptr mEngine; + std::unique_ptr mContext; + + std::string mInputName{"pixels"}; + std::string mOutputName{"latents"}; + rt::Coords mInputShape; + rt::Coords mOutputShape; + nvinfer1::DataType mInputType{nvinfer1::DataType::kHALF}; + nvinfer1::DataType mOutputType{nvinfer1::DataType::kHALF}; + + rt::Tensor mPixelsTensor; + rt::Tensor mLatentsTensor; +}; + +} // namespace rt +} // namespace trt_edgellm diff --git a/cpp/runtime/vlaInferenceRuntime.cpp b/cpp/runtime/vlaInferenceRuntime.cpp new file mode 100644 index 00000000..53657acb --- /dev/null +++ b/cpp/runtime/vlaInferenceRuntime.cpp @@ -0,0 +1,1383 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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 "vlaInferenceRuntime.h" + +#include + +#include "common/bindingNames.h" +#include "common/checkMacros.h" +#include "common/hashUtils.h" +#include "common/logger.h" +#include "common/mathUtils.h" +#include "common/safetensorsUtils.h" +#include "kernels/embeddingKernels/embeddingKernels.h" +#include "kernels/kvCacheUtilKernels/kvCacheUtilsKernels.h" +#include "multimodal/multimodalRunner.h" +#include "multimodal/qwenViTRunner.h" +#include "profiling/metrics.h" +#include "profiling/nvtx_wrapper.h" +#include "profiling/timer.h" +#include "sampler/sampling.h" +#include +#include +#include +#include +#include + +using namespace nvinfer1; + +namespace trt_edgellm +{ + +namespace +{ +std::tuple keySystemPromptWithLoraWeights( + std::string const& systemPrompt, std::string const& loraWeightsName) +{ + return std::make_tuple(systemPrompt, loraWeightsName); +} + +// generateMultimodalIndices is provided by runtime/llmRuntimeUtils.{h,cpp} (rt namespace). + +} // namespace +namespace rt +{ +VlaInferenceRuntime::VlaInferenceRuntime(std::string const& engineDir, std::string const& multimodalEngineDir, + std::unordered_map const& loraWeightsMap, cudaStream_t stream) +{ + // Find the first .engine file in engineDir + // For Qwen3-Omni: export ensures only thinker.engine exists in this directory + std::filesystem::path enginePath; + for (auto const& entry : std::filesystem::directory_iterator(engineDir)) + { + if (entry.path().extension() == ".engine") + { + enginePath = entry.path(); + break; + } + } + if (enginePath.empty()) + { + throw std::runtime_error("No .engine file found in directory: " + engineDir); + } + std::filesystem::path const configPath = std::filesystem::path(engineDir) / "config.json"; + + // Load embedding table from embedding.safetensors + std::filesystem::path const embeddingPath = std::filesystem::path(engineDir) / "embedding.safetensors"; + mEmbedding = loadEmbeddingTable(embeddingPath, stream); + + try + { + mLLMEngineRunner = std::make_unique(enginePath, configPath, loraWeightsMap, stream); + } + catch (std::exception const& e) + { + LOG_ERROR("Failed to initialize LLMEngineRunner: %s", e.what()); + throw std::runtime_error("Failed to initialize LLMEngineRunner: " + std::string(e.what())); + } + LOG_INFO("LLMEngineRunner successfully loaded and initialized llm engine."); + + mEngineConfig = mLLMEngineRunner->getEngineConfig(); + + // Use TopP sampling parameter to reserve max possible workspace size for sampling. + int32_t const defaultTopK{0}; + float const defaultTopP{0.9F}; + trt_edgellm::SamplingParams samplingParams( + mEngineConfig.maxSupportedBatchSize, mEngineConfig.outputVocabSize, 1.0f, defaultTopK, defaultTopP); + int64_t maxSamplingWorkspaceSize = static_cast(trt_edgellm::getTopKtopPSamplingWorkspaceSize( + mEngineConfig.maxSupportedBatchSize, mEngineConfig.outputVocabSize, samplingParams)); + + // Allocate workspace and activation tensors for LLM engine. + try + { + // Use Int8 to indicate byte for workspace. + mSamplingWorkspace = rt::Tensor({maxSamplingWorkspaceSize}, rt::DeviceType::kGPU, DataType::kINT8, + "VlaInferenceRuntime::mSamplingWorkspace"); + mInputIds = rt::Tensor({mEngineConfig.maxSupportedBatchSize, mEngineConfig.maxSupportedInputLength}, + rt::DeviceType::kGPU, DataType::kINT32, "VlaInferenceRuntime::mInputIds"); + mInputsEmbeds = rt::Tensor( + {mEngineConfig.maxSupportedBatchSize, mEngineConfig.maxSupportedInputLength, mEngineConfig.hiddenSize}, + rt::DeviceType::kGPU, DataType::kHALF, "VlaInferenceRuntime::mInputsEmbeds"); + // Allocate deepstack embeddings if needed (one tensor per feature) + if (mEngineConfig.numDeepstackFeatures > 0) + { + mDeepstackEmbeds.resize(mEngineConfig.numDeepstackFeatures); + for (int32_t i = 0; i < mEngineConfig.numDeepstackFeatures; ++i) + { + mDeepstackEmbeds[i] = rt::Tensor({mEngineConfig.maxSupportedBatchSize, + mEngineConfig.maxSupportedInputLength, mEngineConfig.hiddenSize}, + rt::DeviceType::kGPU, DataType::kHALF, + format::fmtstr("VlaInferenceRuntime::mDeepstackEmbeds[%d]", i)); + } + LOG_INFO("Allocated %d deepstack embeds tensors with shape [%d, %d, %d]", + mEngineConfig.numDeepstackFeatures, mEngineConfig.maxSupportedBatchSize, + mEngineConfig.maxSupportedInputLength, mEngineConfig.hiddenSize); + } + mHostPackedInputIds = rt::Tensor({mEngineConfig.maxSupportedBatchSize, mEngineConfig.maxSupportedInputLength}, + rt::DeviceType::kCPU, DataType::kINT32, "VlaInferenceRuntime::mHostPackedInputIds"); + mOutputLogits = rt::Tensor({mEngineConfig.maxSupportedBatchSize, mEngineConfig.outputVocabSize}, + rt::DeviceType::kGPU, DataType::kFLOAT, "VlaInferenceRuntime::mOutputLogits"); + if (mEngineConfig.enableContextEmb || mEngineConfig.enableLmHiddenStates) + { + int32_t const seqOutputDim + = mEngineConfig.enableContextEmb ? mEngineConfig.contextEmbDim : mEngineConfig.hiddenSize; + mOutputContextEmbeds + = rt::Tensor({mEngineConfig.maxSupportedBatchSize, mEngineConfig.maxSupportedInputLength, seqOutputDim}, + rt::DeviceType::kGPU, DataType::kHALF, "VlaInferenceRuntime::mOutputContextEmbeds"); + } + if (mEngineConfig.enablePrefixKVOutputs) + { + check::check(!mEngineConfig.prefixKVOutputShape.empty(), "prefixKVOutputShape must be populated"); + mOutputPrefixK = rt::Tensor(rt::Coords(mEngineConfig.prefixKVOutputShape), rt::DeviceType::kGPU, + DataType::kHALF, "VlaInferenceRuntime::mOutputPrefixK"); + mOutputPrefixV = rt::Tensor(rt::Coords(mEngineConfig.prefixKVOutputShape), rt::DeviceType::kGPU, + DataType::kHALF, "VlaInferenceRuntime::mOutputPrefixV"); + } + mSelectedIndices = rt::Tensor({mEngineConfig.maxSupportedBatchSize, 1}, rt::DeviceType::kGPU, DataType::kINT32, + "VlaInferenceRuntime::mSelectedIndices"); + mHostSelectedTokenIds = rt::Tensor({mEngineConfig.maxSupportedBatchSize}, rt::DeviceType::kCPU, + DataType::kINT32, "VlaInferenceRuntime::mHostSelectedTokenIds"); + mHostContextLengths = rt::Tensor({mEngineConfig.maxSupportedBatchSize}, rt::DeviceType::kCPU, DataType::kINT32, + "VlaInferenceRuntime::mHostContextLengths"); + mHostReuseKVCacheLengths = rt::Tensor({mEngineConfig.maxSupportedBatchSize}, rt::DeviceType::kCPU, + DataType::kINT32, "VlaInferenceRuntime::mHostReuseKVCacheLengths"); + } + catch (std::exception const& e) + { + LOG_ERROR("Failed to allocate workspace and activation tensors for LLM Inference Runtime: %s", e.what()); + throw std::runtime_error( + "Failed to allocate workspace and activation tensors for LLM Inference Runtime: " + std::string(e.what())); + } + + // Setup tokenizer + mTokenizer = std::make_unique(); + LOG_INFO("Start loading tokenizer from model directory: %s", engineDir.c_str()); + if (!mTokenizer->loadFromHF(engineDir)) + { + LOG_ERROR("Failed to load tokenizer from model directory: %s", engineDir.c_str()); + throw std::runtime_error("Failed to load tokenizer from model directory: " + engineDir); + } + mUseCompactPrefixPadding = mTokenizer->getPrefixStrategy() == "pi05_compact_prefix" + || mTokenizer->getPrefixStrategy() == "smolvla_compact_prefix"; + + // Optional: Load vocabulary mapping table if reduced vocabulary is used + if (mEngineConfig.reducedVocabSize > 0) + { + LOG_INFO("Loading vocabulary mapping table for reduced vocab size: %d -> %d", mEngineConfig.reducedVocabSize, + mEngineConfig.vocabSize); + std::filesystem::path const vocabMapPath = std::filesystem::path(engineDir) / binding_names::kVocabMapFileName; + + std::vector vocabMapTensors; + if (!safetensors::loadSafetensors(vocabMapPath, vocabMapTensors, stream)) + { + LOG_ERROR( + "Failed to load %s from model directory: %s", binding_names::kVocabMapFileName, engineDir.c_str()); + throw std::runtime_error("Failed to load " + std::string(binding_names::kVocabMapFileName) + + " from model directory: " + engineDir); + } + + // Check we have exactly one tensor and use it + check::check(vocabMapTensors.size() == 1, + std::string(binding_names::kVocabMapFileName) + " should contain exactly one tensor"); + check::check(vocabMapTensors[0].getShape().getNumDims() == 1, "vocab_map tensor should be 1D"); + check::check(vocabMapTensors[0].getShape()[0] == mEngineConfig.reducedVocabSize, + "vocab_map tensor length should match reduced vocab size"); + mVocabMappingTable = std::move(vocabMapTensors[0]); + LOG_INFO("Vocabulary mapping table successfully loaded."); + } + + // Optional: Setup multimodal engine runners + if (!multimodalEngineDir.empty()) + { + // Multimodal engine directory structure: + // multimodalEngineDir/audio/ - Audio encoder (audio_encoder.engine, config.json) + // multimodalEngineDir/visual/ - Visual encoder (visual.engine, config.json) + // multimodalEngineDir/action/ - Action expert (action.engine) + // + // Note: audio_build and visual_build automatically append /audio and /visual subdirectories. + // Both builders should use the same base --engineDir path. + + // Helper lambda to try loading a runner from a directory + auto tryLoadRunner = [&](std::string const& dir, std::string const& name) -> std::unique_ptr { + try + { + LOG_DEBUG("Attempting to load %s runner from %s", name.c_str(), dir.c_str()); + auto runner = MultimodalRunner::create( + dir, mEngineConfig.maxSupportedBatchSize, mEngineConfig.maxKVCacheCapacity, stream); + LOG_INFO("%s runner successfully initialized", name.c_str()); + return runner; + } + catch (std::exception const& e) + { + LOG_DEBUG("Failed to load %s runner from %s: %s", name.c_str(), dir.c_str(), e.what()); + return nullptr; + } + }; + + // Try to load audio runner from multimodalEngineDir/audio + mAudioRunner = tryLoadRunner(multimodalEngineDir + "/audio", "Audio"); + + // Try to load visual runner from multimodalEngineDir/visual (with fallback to root for pure visual models) + mVisionRunner = tryLoadRunner(multimodalEngineDir + "/visual", "Visual"); + if (!mVisionRunner) + { + mVisionRunner = tryLoadRunner(multimodalEngineDir, "Vision"); + } + + // At least one runner must be available + if (!mAudioRunner && !mVisionRunner) + { + throw std::runtime_error("No valid multimodal engine found in " + multimodalEngineDir); + } + + try + { + std::string const actionCtxDir = multimodalEngineDir + "/action_context"; + LOG_INFO("Attempting to load action context runner from %s", actionCtxDir.c_str()); + mActionContextRunner = std::make_unique(actionCtxDir, stream); + LOG_INFO("Action context runner loaded from %s", actionCtxDir.c_str()); + } + catch (std::exception const& e) + { + LOG_INFO("Failed to load action context runner from %s: %s", + (multimodalEngineDir + "/action_context").c_str(), e.what()); + } + + // Try to load action expert from multimodalEngineDir/action + try + { + std::string const actionDir = multimodalEngineDir + "/action"; + LOG_INFO("Attempting to load Action runner from %s", actionDir.c_str()); + mActionRunner + = std::make_unique(actionDir, stream, mLLMEngineRunner->getLinearKVCache().getConfig()); + LOG_INFO("Action runner loaded (handoff=%s, rollout=%s).", + mActionRunner->getContextHandoff() == ActionContextHandoff::PREFIX_KV ? "prefix_kv" : "context_tensor", + mActionRunner->getRolloutMode() == ActionRolloutMode::FLOW_MATCHING ? "flow_matching" : "velocity"); + } + catch (std::exception const& e) + { + LOG_INFO("Failed to load Action runner from %s: %s", (multimodalEngineDir + "/action").c_str(), e.what()); + } + + // Prefix-KV action engines must match the LM KV cache capacity. + if (mActionRunner && mActionRunner->getContextHandoff() == ActionContextHandoff::PREFIX_KV) + { + int32_t const actionMaxKVCacheCapacity = mActionRunner->getMaxKVCacheCapacity(); + int32_t const llmMaxKVCacheCapacity = mEngineConfig.maxKVCacheCapacity; + if (actionMaxKVCacheCapacity != llmMaxKVCacheCapacity) + { + throw std::runtime_error(format::fmtstr( + "Action engine max_kv_cache_capacity (%d) does not match LLM engine max_kv_cache_capacity (%d). " + "Re-export and rebuild the action engine with --max_kv_cache_capacity=%d to match the LLM engine.", + actionMaxKVCacheCapacity, llmMaxKVCacheCapacity, llmMaxKVCacheCapacity)); + } + } + } + + // Setup shared execution context memory for LLM and multimodal engines. + // All engines execute serially (not concurrently), so they can share a single buffer + // sized to the maximum requirement among all engines. + int64_t const llmContextMemorySize = mLLMEngineRunner->getRequiredContextMemorySize(); + int64_t const visionContextMemorySize = mVisionRunner ? mVisionRunner->getRequiredContextMemorySize() : 0; + int64_t const audioContextMemorySize = mAudioRunner ? mAudioRunner->getRequiredContextMemorySize() : 0; + int64_t const actionContextRunnerMemorySize + = mActionContextRunner ? mActionContextRunner->getRequiredContextMemorySize() : 0; + int64_t const actionContextMemorySize = mActionRunner ? mActionRunner->getRequiredContextMemorySize() : 0; + int64_t const sharedContextMemorySize + = std::max({llmContextMemorySize, visionContextMemorySize, audioContextMemorySize, actionContextMemorySize}); + mSharedExecContextMemory = rt::Tensor({sharedContextMemorySize}, rt::DeviceType::kGPU, nvinfer1::DataType::kUINT8, + "VlaInferenceRuntime::mSharedExecContextMemory"); + mLLMEngineRunner->setContextMemory(mSharedExecContextMemory); + if (mVisionRunner) + { + mVisionRunner->setContextMemory(mSharedExecContextMemory); + } + if (mAudioRunner) + { + mAudioRunner->setContextMemory(mSharedExecContextMemory); + } + if (mActionContextRunner) + { + mActionContextRunner->setContextMemory(mSharedExecContextMemory); + } + if (mActionRunner) + { + mActionRunner->setContextMemory(mSharedExecContextMemory); + } + LOG_INFO( + "Setup shared execution context memory: %zu bytes (llm requires: %zu, vision requires: %zu, audio " + "requires: %zu, action_context requires: %zu, action requires: %zu)", + static_cast(sharedContextMemorySize), static_cast(llmContextMemorySize), + static_cast(visionContextMemorySize), static_cast(audioContextMemorySize), + static_cast(actionContextRunnerMemorySize), static_cast(actionContextMemorySize)); +} + +void VlaInferenceRuntime::setActionNoiseSeed(int32_t seed) noexcept +{ + if (mActionRunner) + { + mActionRunner->setNoiseSeed(seed); + } +} + +bool VlaInferenceRuntime::examineRequest(LLMGenerationRequest const& request) noexcept +{ + int32_t const activeBatchSize = static_cast(request.requests.size()); + + if (activeBatchSize == 0) + { + LOG_ERROR("VlaInferenceRuntime(): The request is empty with no requests supplied."); + return false; + } + + if (activeBatchSize > mEngineConfig.maxSupportedBatchSize) + { + LOG_ERROR("VlaInferenceRuntime(): The batched request size (%d) exceeds the max supported batch size (%d).", + activeBatchSize, mEngineConfig.maxSupportedBatchSize); + return false; + } + + for (auto const& request : request.requests) + { + if (request.messages.empty()) + { + LOG_ERROR( + "There is an empty request in the batch. 'messages' must be provided. " + "Skip this batch of requests. Please check the input data contents."); + return false; + } + } + + return true; +} + +bool VlaInferenceRuntime::setUpForPrefillExecution(std::vector> const& batchedInputIds, + std::vector const& systemPrompts, std::string const& loraWeightsName, cudaStream_t stream) +{ + NVTX_SCOPED_RANGE(nvtx_setup, "SETUP_PREFILL_EXECUTION", nvtx_colors::PALE_GREEN); + + std::vector> processedInputIds; + std::vector processedIdsLengths; + int32_t const activeBatchSize = static_cast(batchedInputIds.size()); + + rt::LinearKVCache& linearKVCache = mLLMEngineRunner->getLinearKVCache(); + rt::Tensor kvCacheBuffer = linearKVCache.getKVCacheBuffer(); + + // Record the length of the reused KVCache for each sequence using pre-allocated tensor + check::check(mHostReuseKVCacheLengths.reshape({activeBatchSize}), "Tensor reshape failed"); + int32_t* reuseKVCacheLengthsData = mHostReuseKVCacheLengths.dataPointer(); + + // Search if the system prompt has been cached. If there are cached system prompts, insert + // the pre-computed KVCache and remove the contents from inputIds. + for (int32_t i = 0; i < activeBatchSize; ++i) + { + auto const promptKey = keySystemPromptWithLoraWeights(systemPrompts[i], loraWeightsName); + if (mSystemPromptKVCache.find(promptKey) != mSystemPromptKVCache.end()) + { + auto& precachedKVCache = mSystemPromptKVCache[promptKey]; + auto const& kvCacheContent = precachedKVCache.kvCacheContent; + kernel::instantiateKVCacheFromTensor(kvCacheBuffer, kvCacheContent, i, stream); + auto reuseLength = math::cast(kvCacheContent.getShape()[3]); + check::check( + reuseLength < batchedInputIds[i].size(), "The reuse length shall not exceed the input length."); + processedInputIds.emplace_back(batchedInputIds[i].begin() + reuseLength, batchedInputIds[i].end()); + processedIdsLengths.emplace_back(math::cast(batchedInputIds[i].size() - reuseLength)); + reuseKVCacheLengthsData[i] = math::cast(reuseLength); + // If the system prompt is not well designed, the boundary of the inputIDs could be mis-aligned. + bool const matchIds = std::equal(precachedKVCache.tokenizedPrompt.begin(), + precachedKVCache.tokenizedPrompt.end(), batchedInputIds[i].begin()); + if (!matchIds) + { + LOG_WARNING( + "VlaInferenceRuntime(): Though system prompt strings are matched, token_ids are not perfectly " + "aligned. " + "This may generate incorrect result, please check your system prompt design."); + } + } + else + { + processedInputIds.emplace_back(batchedInputIds[i]); + processedIdsLengths.emplace_back(static_cast(batchedInputIds[i].size())); + reuseKVCacheLengthsData[i] = 0; + } + } + + // Pack inputIds, instantiate input data for prefill step, and reset the KVCache state. + int32_t const maxInputLength = *std::max_element(processedIdsLengths.begin(), processedIdsLengths.end()); + if (maxInputLength > mEngineConfig.maxSupportedInputLength) + { + LOG_ERROR( + "VlaInferenceRuntime(): The max input length (%d) exceeds the max supported input length (%d) of the LLM " + "Engine.", + maxInputLength, mEngineConfig.maxSupportedInputLength); + return false; + } + + // Pad each batch to engine max length for static-shape VLA engines (pi05_compact_prefix), + // otherwise pad only to the batch max valid length. + int32_t const packedInputLength = mUseCompactPrefixPadding ? mEngineConfig.maxSupportedInputLength : maxInputLength; + if (maxInputLength > packedInputLength) + { + LOG_ERROR( + "VlaInferenceRuntime(): The max input length (%d) exceeds the packed input length (%d) of the LLM " + "Engine.", + maxInputLength, packedInputLength); + return false; + } + check::check(mHostPackedInputIds.reshape({activeBatchSize, packedInputLength}), "Tensor reshape failed"); + int32_t* packedInputIdsData = mHostPackedInputIds.dataPointer(); + std::fill(packedInputIdsData, packedInputIdsData + activeBatchSize * packedInputLength, mTokenizer->getPadId()); + + for (int32_t i = 0; i < activeBatchSize; ++i) + { + // Pad each sequence to the max length of this batch. + // TODO: Implement remove input padding for better efficiency until multi-batch. + std::copy(processedInputIds[i].begin(), processedInputIds[i].end(), packedInputIdsData + i * packedInputLength); + } + + linearKVCache.resetForNewSequences(mHostReuseKVCacheLengths, stream); + + // For each recurrent layer and each batch element, either restore the cached + // recurrent/conv state if the cache hit, or zero the state. + if (mEngineConfig.numLinearAttnLayers > 0) + { + rt::LinearKVCache& kvCache = mLLMEngineRunner->getLinearKVCache(); + rt::LinearKVCache::CacheConfig const& cacheConfig = kvCache.getConfig(); + size_t const recurrentElemSize = rt::utils::getTypeSize(cacheConfig.recurrentStateType); + size_t const convElemSize = rt::utils::getTypeSize(cacheConfig.convStateType); + size_t const recurrentBatchBytes = static_cast(cacheConfig.recurrentStateNumHeads + * cacheConfig.recurrentStateHeadDim * cacheConfig.recurrentStateSize) + * recurrentElemSize; + size_t const convBatchBytes = static_cast(cacheConfig.convDim * cacheConfig.convKernel) * convElemSize; + + for (int32_t layer = 0; layer < mEngineConfig.numLinearAttnLayers; ++layer) + { + rt::Tensor recurrentLayer = kvCache.getRecurrentStateForLayer(layer); + rt::Tensor convLayer = kvCache.getConvStateForLayer(layer); + + for (int32_t i = 0; i < activeBatchSize; ++i) + { + auto* recurrentDst = static_cast(recurrentLayer.rawPointer()) + i * recurrentBatchBytes; + auto* convDst = static_cast(convLayer.rawPointer()) + i * convBatchBytes; + + auto const promptKey = keySystemPromptWithLoraWeights(systemPrompts[i], loraWeightsName); + auto it = mSystemPromptKVCache.find(promptKey); + bool const hasCache = (it != mSystemPromptKVCache.end()); + + if (hasCache && layer < static_cast(it->second.recurrentStateContents.size())) + { + CUDA_CHECK(cudaMemcpyAsync(recurrentDst, it->second.recurrentStateContents[layer].rawPointer(), + recurrentBatchBytes, cudaMemcpyDeviceToDevice, stream)); + } + else + { + CUDA_CHECK(cudaMemsetAsync(recurrentDst, 0, recurrentBatchBytes, stream)); + } + + if (hasCache && layer < static_cast(it->second.convStateContents.size())) + { + CUDA_CHECK(cudaMemcpyAsync(convDst, it->second.convStateContents[layer].rawPointer(), + convBatchBytes, cudaMemcpyDeviceToDevice, stream)); + } + else + { + CUDA_CHECK(cudaMemsetAsync(convDst, 0, convBatchBytes, stream)); + } + } + } + } + + check::check(mInputIds.reshape({activeBatchSize, packedInputLength}), "Tensor reshape failed"); + check::check(mHostContextLengths.reshape({activeBatchSize}), "Tensor reshape failed"); + if (mEngineConfig.hasLogitsOutput) + { + check::check(mOutputLogits.reshape({activeBatchSize, mEngineConfig.outputVocabSize}), "Tensor reshape failed"); + } + if (mEngineConfig.enableContextEmb || mEngineConfig.enableLmHiddenStates) + { + int32_t const seqOutputDim + = mEngineConfig.enableContextEmb ? mEngineConfig.contextEmbDim : mEngineConfig.hiddenSize; + check::check( + mOutputContextEmbeds.reshape({activeBatchSize, packedInputLength, seqOutputDim}), "Tensor reshape failed"); + } + if (mEngineConfig.enablePrefixKVOutputs) + { + rt::Coords prefixShape(mEngineConfig.prefixKVOutputShape); + if (prefixShape.getNumDims() > 0) + { + prefixShape[1] = activeBatchSize; + } + check::check(mOutputPrefixK.reshape(prefixShape), "Tensor reshape failed"); + check::check(mOutputPrefixV.reshape(prefixShape), "Tensor reshape failed"); + } + + CUDA_CHECK(cudaMemcpyAsync(mInputIds.rawPointer(), mHostPackedInputIds.rawPointer(), + activeBatchSize * packedInputLength * sizeof(int32_t), cudaMemcpyHostToDevice, stream)); + memcpy(mHostContextLengths.dataPointer(), processedIdsLengths.data(), activeBatchSize * sizeof(int32_t)); + + if (mEngineConfig.maxSupportedLoraRank > 0 && !mLLMEngineRunner->switchLoraWeights(loraWeightsName)) + { + LOG_ERROR("Failed to switch LoRA weights to %s", loraWeightsName.c_str()); + return false; + } + + return true; +} + +bool VlaInferenceRuntime::handleRequest( + LLMGenerationRequest const& request, LLMGenerationResponse& response, cudaStream_t stream) +{ + auto const e2eStart = std::chrono::steady_clock::now(); + + std::vector> batchedInputIds; + std::vector batchSystemPrompts; + std::string loraWeightsName = request.loraWeightsName; + + if (!examineRequest(request)) + { + LOG_ERROR("VlaInferenceRuntime(): Input request examination failed. This request cannot be handled."); + return false; + } + + int32_t activeBatchSize = static_cast(request.requests.size()); + + // Apply chat template, extract system prompts, and optionally save KVCache + request.formattedRequests.resize(activeBatchSize); + batchSystemPrompts.reserve(activeBatchSize); + + for (int32_t i = 0; i < activeBatchSize; ++i) + { + // Apply chat template + mTokenizer->applyChatTemplate(request.requests[i], request.formattedRequests[i], request.applyChatTemplate, + request.addGenerationPrompt, request.enableThinking); + + // Extract system prompt + batchSystemPrompts.emplace_back(request.formattedRequests[i].formattedSystemPrompt); + + // Save KVCache if requested + if (request.saveSystemPromptKVCache) + { + if (mVisionRunner) + { + mVisionRunner->preprocessSystemPrompt( + batchSystemPrompts[i], mTokenizer.get(), mLLMEngineRunner->getRopeCosSinCacheTensor(), stream); + } + else if (mAudioRunner) + { + mAudioRunner->preprocessSystemPrompt( + batchSystemPrompts[i], mTokenizer.get(), mLLMEngineRunner->getRopeCosSinCacheTensor(), stream); + } + bool const saveCacheStatus = genAndSaveSystemPromptKVCache(batchSystemPrompts[i], loraWeightsName, stream); + if (!saveCacheStatus) + { + LOG_WARNING( + "Failed to save system prompt KVCache. Continue to handle the request without saving the system " + "prompt KVCache."); + } + } + } + + // Preprocess user prompts and encode them. + // Check if request has audio or vision inputs + bool hasAudio = std::any_of( + request.requests.begin(), request.requests.end(), [](auto const& req) { return !req.audioBuffers.empty(); }); + bool hasVision = std::any_of( + request.requests.begin(), request.requests.end(), [](auto const& req) { return !req.imageBuffers.empty(); }); + bool const hasTrajectoryHistory = std::any_of(request.requests.begin(), request.requests.end(), + [](auto const& req) { return req.pastTrajectory.has_value(); }); + + int32_t const actionBatchSize = (request.actionBatchSize > 0) ? request.actionBatchSize : activeBatchSize; + + bool const runPrefixKvAction = mActionRunner != nullptr + && mActionRunner->getContextHandoff() == ActionContextHandoff::PREFIX_KV && hasTrajectoryHistory; + bool const runPi05VelocityAction = mActionRunner != nullptr + && mActionRunner->getContextHandoff() == ActionContextHandoff::CONTEXT_TENSOR && hasVision + && mEngineConfig.enablePrefixKVOutputs && mActionContextRunner == nullptr; + bool const runContextTensorAction = mActionRunner != nullptr + && mActionRunner->getContextHandoff() == ActionContextHandoff::CONTEXT_TENSOR && hasVision + && !runPi05VelocityAction + && ((mEngineConfig.enableContextEmb && mActionContextRunner == nullptr) + || (mActionContextRunner != nullptr && mEngineConfig.enableLmHiddenStates)); + + if ((hasAudio && mAudioRunner) || (hasVision && mVisionRunner)) + { + // Mark multimodal preprocessing and inference for NVTX profiling + NVTX_SCOPED_RANGE(nvtx_multimodal, "MULTIMODAL_PROCESSING", nvtx_colors::ORANGE); + + // Process audio inputs (if present) + if (hasAudio && mAudioRunner) + { + LOG_INFO("Processing audio inputs"); + if (!mAudioRunner->preprocess( + request, batchedInputIds, mTokenizer.get(), mLLMEngineRunner->getRopeCosSinCacheTensor(), stream)) + { + LOG_ERROR("VlaInferenceRuntime(): Audio preprocessing failed. This request cannot be handled."); + return false; + } + + if (!mAudioRunner->infer(stream)) + { + LOG_ERROR("VlaInferenceRuntime(): Audio inference failed. This request cannot be handled."); + return false; + } + } + + // Process vision inputs (if present) + if (hasVision && mVisionRunner) + { + LOG_INFO("Processing vision inputs"); + auto const vitStart = std::chrono::steady_clock::now(); + if (!mVisionRunner->preprocess( + request, batchedInputIds, mTokenizer.get(), mLLMEngineRunner->getRopeCosSinCacheTensor(), stream)) + { + LOG_ERROR("VlaInferenceRuntime(): Vision preprocessing failed. This request cannot be handled."); + return false; + } + + if (!mVisionRunner->infer(stream)) + { + LOG_ERROR("VlaInferenceRuntime(): Vision inference failed. This request cannot be handled."); + return false; + } + auto const vitEnd = std::chrono::steady_clock::now(); + double const vitMs = std::chrono::duration(vitEnd - vitStart).count(); + LOG_INFO("Stage timings - ViT: %.3f ms", vitMs); + } + } + else + { + // Pure text mode: directly tokenize + batchedInputIds.reserve(activeBatchSize); + for (int32_t i = 0; i < activeBatchSize; ++i) + { + batchedInputIds.emplace_back( + mTokenizer->encode(request.formattedRequests[i].formattedCompleteRequest, true)); + if (batchedInputIds[i].empty()) + { + LOG_ERROR("Failed to encode input text for request %d in batch", i); + return false; + } + } + } + + if (runPrefixKvAction || runContextTensorAction || runPi05VelocityAction) + { + LOG_INFO("Preprocessing action inputs (LLM batch=%d, action batch=%d, handoff=%s)", activeBatchSize, + actionBatchSize, + runPrefixKvAction ? "prefix_kv" : (runPi05VelocityAction ? "pi05_prefix_kv" : "context_tensor")); + if (!mActionRunner->preprocess(request, batchedInputIds, mTokenizer.get())) + { + LOG_ERROR("VlaInferenceRuntime(): Action preprocessing failed. This request cannot be handled."); + return false; + } + } + + // Conduct the preparation work to handle a new set of sequences, including inputIds packing, input/output tensor + // preparation, reset the KVCache state, and apply reused prefix KVCache if available. + if (!setUpForPrefillExecution(batchedInputIds, batchSystemPrompts, loraWeightsName, stream)) + { + LOG_ERROR("VlaInferenceRuntime(): Prefill execution setup failed. This request cannot be handled."); + return false; + } + + // Record context information for performance tracking + auto tokenCount = calculateTokenCounts(batchedInputIds, batchSystemPrompts, loraWeightsName); + + int32_t const* contextLengthsData = mHostContextLengths.dataPointer(); + int32_t actualContextLength{0}; + for (int32_t i = 0; i < activeBatchSize; ++i) + { + actualContextLength = std::max(actualContextLength, contextLengthsData[i]); + } + int32_t maxGenerationLength = request.maxGenerateLength; + if (actualContextLength + maxGenerationLength > mEngineConfig.maxKVCacheCapacity) + { + maxGenerationLength = mEngineConfig.maxKVCacheCapacity - actualContextLength; + LOG_WARNING( + "The requested input length (%d) + max generation length (%d) = %d exceeds the max KV " + "cache capacity (%d). Reduce the generation length to %d to avoid the truncation of the generated tokens.", + actualContextLength, request.maxGenerateLength, actualContextLength + request.maxGenerateLength, + mEngineConfig.maxKVCacheCapacity, maxGenerationLength); + } + + // Set up data structures to store the generated results during decoding. + // Also set up sampling parameters and sampling lambda function. + int32_t unFinishedBatchNum = activeBatchSize; + int32_t generationIter{0}; + std::vector> outputIds(activeBatchSize); + std::vector finishedStates(activeBatchSize, false); + check::check(mSelectedIndices.reshape({activeBatchSize, 1}), "Tensor reshape failed"); + check::check(mHostSelectedTokenIds.reshape({activeBatchSize}), "Tensor reshape failed"); + int32_t* hostSelectedTokenIdsData = mHostSelectedTokenIds.dataPointer(); + + // Used for prefix-KV trajectory models (e.g. Alpamayo): stop decode after traj marker token. + int32_t trajFutureStartId = 0; + if (runPrefixKvAction) + { + trajFutureStartId = static_cast(mTokenizer->getTokenId("<|traj_future_start|>")); + } + SamplingParams params( + activeBatchSize, mEngineConfig.outputVocabSize, request.temperature, request.topK, request.topP); + auto sampleTokens = [&]() { + trt_edgellm::topKtopPSamplingFromLogits(mOutputLogits, mSelectedIndices, params, mSamplingWorkspace, stream); + // Apply vocabulary mapping if reduced vocabulary is used + if (mEngineConfig.reducedVocabSize > 0) + { + trt_edgellm::mapReducedVocabToFullVocab(mSelectedIndices, mVocabMappingTable, stream); + } + CUDA_CHECK(cudaMemcpyAsync(mHostSelectedTokenIds.rawPointer(), mSelectedIndices.rawPointer(), + activeBatchSize * sizeof(int32_t), cudaMemcpyDeviceToHost, stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + for (int32_t i = 0; i < activeBatchSize; ++i) + { + if (!finishedStates[i]) + { + outputIds[i].push_back(hostSelectedTokenIdsData[i]); + + // Prefix-KV trajectory models stop one token after traj_future_start so KV includes the marker. + if (runPrefixKvAction) + { + if (outputIds[i].size() > 1 && trajFutureStartId >= 0) + { + finishedStates[i] = outputIds[i][outputIds[i].size() - 2] == trajFutureStartId; + } + } + else + { + finishedStates[i] = hostSelectedTokenIdsData[i] == mTokenizer->getEosId(); + } + + if (finishedStates[i]) + { + unFinishedBatchNum--; + } + } + } + ++generationIter; + }; + + // Perform embedding lookup for prefill + int32_t const prefillSequenceLength = mInputIds.getShape()[1]; + check::check(mInputsEmbeds.reshape({activeBatchSize, prefillSequenceLength, mEngineConfig.hiddenSize}), + "Tensor reshape failed"); + + // Get embeddings from independent runners + rt::OptionalInputTensor visionEmbeddings + = mVisionRunner ? std::optional{std::ref(mVisionRunner->getOutputEmbedding())} : std::nullopt; + rt::OptionalInputTensor audioEmbeddings + = mAudioRunner ? std::optional{std::ref(mAudioRunner->getOutputEmbedding())} : std::nullopt; + + if (audioEmbeddings.has_value()) + { + // Audio present: use embeddingLookupMultimodal (handles audio and/or vision) + auto const inputShape = mInputIds.getShape(); + size_t const inputSizeBytes = inputShape.volume() * sizeof(int32_t); + rt::Tensor inputIdsCPU(inputShape, rt::DeviceType::kCPU, mInputIds.getDataType()); + CUDA_CHECK( + cudaMemcpy(inputIdsCPU.rawPointer(), mInputIds.rawPointer(), inputSizeBytes, cudaMemcpyDeviceToHost)); + + std::optional audioTokenId + = (mEngineConfig.audioTokenId != 0) ? std::optional{mEngineConfig.audioTokenId} : std::nullopt; + std::optional imageTokenId + = (mEngineConfig.imageTokenId != 0) ? std::optional{mEngineConfig.imageTokenId} : std::nullopt; + rt::Tensor multimodalIndicesCPU + = generateMultimodalIndices(inputIdsCPU, audioTokenId, imageTokenId, mEngineConfig.vocabSize); + + auto const indicesShape = multimodalIndicesCPU.getShape(); + size_t const indicesSizeBytes = indicesShape.volume() * sizeof(int32_t); + mMultimodalIndices = rt::Tensor(indicesShape, rt::DeviceType::kGPU, multimodalIndicesCPU.getDataType()); + CUDA_CHECK(cudaMemcpy(mMultimodalIndices.rawPointer(), multimodalIndicesCPU.rawPointer(), indicesSizeBytes, + cudaMemcpyHostToDevice)); + + kernel::embeddingLookupMultimodal(mInputIds, mEmbedding.table, mEmbedding.scalesAsOptional(), + std::optional{std::ref(mMultimodalIndices)}, imageTokenId, visionEmbeddings, audioTokenId, audioEmbeddings, + mInputsEmbeds, stream); + } + else if (visionEmbeddings.has_value()) + { + // Legacy vision path (Qwen2.5-VL, InternVL: imageTokenId >= vocabSize or not set) + rt::Tensor const& imageEmbedsTensor = visionEmbeddings.value().get(); + kernel::embeddingLookupWithImageInsertion( + mInputIds, mEmbedding.table, mEmbedding.scalesAsOptional(), imageEmbedsTensor, mInputsEmbeds, stream); + } + else + { + // Standard embedding lookup (pure text) + kernel::embeddingLookup(mInputIds, mEmbedding.table, mEmbedding.scalesAsOptional(), mInputsEmbeds, stream); + } + + // Process deepstack features: perform embedding assembly if vision runner is available + // Note: Deepstack features are only provided by VisionRunner, not Qwen3OmniAudioRunner + rt::OptionalInputTensors deepstackEmbeds{}; + if (mEngineConfig.numDeepstackFeatures > 0 && mVisionRunner) + { + rt::OptionalInputTensors deepstackFeatures = mVisionRunner->getDeepstackFeatures(); + + // Prepare multimodal indices for deepstack assembly (needed when imageTokenId < vocabSize) + rt::OptionalInputTensor deepstackMultimodalIndices{std::nullopt}; + if (mMultimodalIndices.getShape().volume() > 0) + { + deepstackMultimodalIndices = std::ref(mMultimodalIndices); + } + + for (int32_t idx = 0; idx < static_cast(deepstackFeatures.size()); ++idx) + { + rt::Tensor const& featureTensor = deepstackFeatures[idx].get(); + + // Reshape and perform embedding assembly for this feature + check::check( + mDeepstackEmbeds[idx].reshape({activeBatchSize, prefillSequenceLength, mEngineConfig.hiddenSize}), + "Tensor reshape failed"); + kernel::assembleDeepstackEmbedding(mInputIds, featureTensor, mEngineConfig.vocabSize, mDeepstackEmbeds[idx], + stream, mEngineConfig.imageTokenId, deepstackMultimodalIndices); + + // Add to output vector (engine will bind by index) + deepstackEmbeds.push_back(std::ref(mDeepstackEmbeds[idx])); + } + } + + // Profile all sampling operations as one stage + // Prefill profiling session + auto const llmPrefillStart = std::chrono::steady_clock::now(); + { + TIME_STAGE(metrics::StageNames::kLLM_PREFILL, stream); + // Enhanced NVTX range with detailed information + NVTX_SCOPED_RANGE(nvtx_prefill, + ("LLM_PREFILL[BS=" + std::to_string(activeBatchSize) + + ",Reused=" + std::to_string(tokenCount.totalReusedTokens) + + ",Computed=" + std::to_string(tokenCount.totalComputedTokens) + "]") + .c_str(), + nvtx_colors::BLUE); + + rt::OptionalOutputTensor outputContextEmbeds{std::nullopt}; + rt::OptionalOutputTensor outputPrefixK{std::nullopt}; + rt::OptionalOutputTensor outputPrefixV{std::nullopt}; + if (mEngineConfig.enablePrefixKVOutputs) + { + outputPrefixK = std::ref(mOutputPrefixK); + outputPrefixV = std::ref(mOutputPrefixV); + } + if (mEngineConfig.enableContextEmb || mEngineConfig.enableLmHiddenStates) + { + outputContextEmbeds = std::ref(mOutputContextEmbeds); + } + bool prefillStatus + = mLLMEngineRunner->executePrefillStep(mInputsEmbeds, mHostContextLengths, deepstackEmbeds, mOutputLogits, + rt::OptionalOutputTensor{std::nullopt}, stream, outputContextEmbeds, outputPrefixK, outputPrefixV); + if (!prefillStatus) + { + LOG_ERROR( + "VlaInferenceRuntime(): Failed to execute prefill step. Cannot generate the KVCache for this prompt."); + return false; + } + + if (mEngineConfig.hasLogitsOutput) + { + sampleTokens(); + } + } + auto const llmPrefillEnd = std::chrono::steady_clock::now(); + double const llmPrefillMs = std::chrono::duration(llmPrefillEnd - llmPrefillStart).count(); + LOG_INFO("Stage timings - LLM Prefill: %.3f ms", llmPrefillMs); + + // Record prefill metrics + mPrefillMetrics.recordRun(tokenCount.totalReusedTokens, tokenCount.totalComputedTokens); + + // Reshape for decoding step + check::check(mInputsEmbeds.reshape({activeBatchSize, 1, mEngineConfig.hiddenSize}), "Tensor reshape failed"); + + // Profile entire generation phase like benchmark profiler + auto const llmGenStart = std::chrono::steady_clock::now(); + { + TIME_STAGE(metrics::StageNames::kLLM_GENERATION, stream); + // Enhanced NVTX range with batch size + NVTX_SCOPED_RANGE(nvtx_generation, + ("LLM_GENERATION[BS=" + std::to_string(activeBatchSize) + ",MaxLen=" + std::to_string(maxGenerationLength) + + "]") + .c_str(), + nvtx_colors::GREEN); + + while (unFinishedBatchNum > 0 && generationIter < maxGenerationLength) + { + // Mark each decoding iteration with detailed info + NVTX_SCOPED_RANGE(iter_range, + ("Decode_Iter[" + std::to_string(generationIter) + "/" + std::to_string(maxGenerationLength) + + ",Active=" + std::to_string(unFinishedBatchNum) + "]") + .c_str(), + nvtx_colors::LIGHT_GREEN); + + // Perform embedding lookup for the selected token indices (decode only has text, no images) + kernel::embeddingLookup( + mSelectedIndices, mEmbedding.table, mEmbedding.scalesAsOptional(), mInputsEmbeds, stream); + + // Use the embedded tokens as input for the decoding step. + // No hidden states output needed for standard LLM decoding. + rt::OptionalOutputTensor const outputHiddenStates{std::nullopt}; + bool decodingStatus = mLLMEngineRunner->executeVanillaDecodingStep( + mInputsEmbeds, mOutputLogits, outputHiddenStates, stream); + if (!decodingStatus) + { + LOG_ERROR("VlaInferenceRuntime(): Failed to execute decoding step."); + return false; + } + + sampleTokens(); + } + } + auto const llmGenEnd = std::chrono::steady_clock::now(); + double const llmGenMs = std::chrono::duration(llmGenEnd - llmGenStart).count(); + LOG_INFO("Stage timings - LLM Generation: %.3f ms", llmGenMs); + + // Record generation and sampling metrics + int32_t totalGeneratedTokens = 0; + for (int32_t i = 0; i < activeBatchSize; ++i) + { + totalGeneratedTokens += static_cast(outputIds[i].size() - 1); + } + + if (totalGeneratedTokens > 0) + { + mGenerationMetrics.recordRun(totalGeneratedTokens); + } + + // Clean the response field and fill the generated outputIds and decoded texts. + response.outputIds.clear(); + response.outputTexts.clear(); + response.outputTrajectories.clear(); + response.outputActions.clear(); + for (int32_t i = 0; i < activeBatchSize; ++i) + { + response.outputIds.emplace_back(outputIds[i]); + response.outputTexts.emplace_back(mTokenizer->decode(outputIds[i], true)); + } + if (actionBatchSize > activeBatchSize && !response.outputIds.empty()) + { + std::vector const llmOutputIds = response.outputIds[0]; + std::string const llmOutputText = response.outputTexts[0]; + response.outputIds.assign(actionBatchSize, llmOutputIds); + response.outputTexts.assign(actionBatchSize, llmOutputText); + activeBatchSize = actionBatchSize; + } + + if (runPrefixKvAction) + { + if (!mVisionRunner) + { + LOG_ERROR("Prefix-KV action runner requires a vision runner for MRoPE rope deltas."); + return false; + } + + multimodal::ModelType const visionType = mVisionRunner->getModelType(); + bool const isQwen3ViT = visionType == multimodal::ModelType::QWEN3_VL; + if (!isQwen3ViT) + { + LOG_ERROR( + "Prefix-KV action runner requires a Qwen3-VL vision runner but a different vision runner is loaded."); + return false; + } + + auto* qwenVision = static_cast(mVisionRunner.get()); + std::vector const& ropeDeltasFromVit = qwenVision->getMropeRopeDeltasPerBatch(); + std::vector ropeDeltasBroadcast; + std::vector const* ropeDeltasPtr = &ropeDeltasFromVit; + if (static_cast(ropeDeltasFromVit.size()) < activeBatchSize) + { + int64_t const ropeDelta0 = ropeDeltasFromVit.empty() ? int64_t{0} : ropeDeltasFromVit[0]; + ropeDeltasBroadcast.assign(activeBatchSize, ropeDelta0); + ropeDeltasPtr = &ropeDeltasBroadcast; + } + std::vector const& ropeDeltas = *ropeDeltasPtr; + response.outputTrajectories.resize(activeBatchSize); + rt::LinearKVCache& kvcache = mLLMEngineRunner->getLinearKVCache(); + auto const diffusorStart = std::chrono::steady_clock::now(); + std::vector> trajectories + = mActionRunner->sampleTrajectory(stream, activeBatchSize, kvcache, ropeDeltas); + auto const diffusorEnd = std::chrono::steady_clock::now(); + double const diffusorMs = std::chrono::duration(diffusorEnd - diffusorStart).count(); + LOG_INFO("Stage timings - Diffusor: %.3f ms", diffusorMs); + if (trajectories.size() != static_cast(activeBatchSize)) + { + LOG_ERROR("VlaInferenceRuntime(): prefix-KV action sampling failed."); + return false; + } + for (size_t i = 0; i < trajectories.size() && i < static_cast(activeBatchSize); ++i) + { + if (!trajectories[i].empty()) + { + response.outputTrajectories[i] = std::move(trajectories[i]); + } + } + } + else if (runPi05VelocityAction) + { + auto const diffusorStart = std::chrono::steady_clock::now(); + if (!mActionRunner->wireStaticInputs(request, stream)) + { + LOG_ERROR("VlaInferenceRuntime(): failed to wire static action inputs."); + return false; + } + + int32_t prefixValidLen{0}; + int32_t const* contextLengthsData = mHostContextLengths.dataPointer(); + for (int32_t i = 0; i < activeBatchSize; ++i) + { + prefixValidLen = std::max(prefixValidLen, contextLengthsData[i]); + } + if (!mActionRunner->preparePi05SuffixInputs(activeBatchSize, prefixValidLen, stream)) + { + LOG_ERROR("VlaInferenceRuntime(): failed to prepare PI0.5 suffix action inputs."); + return false; + } + + std::vector languageOutputs(3, nullptr); + languageOutputs[0] = &mOutputContextEmbeds; + languageOutputs[1] = &mOutputPrefixK; + languageOutputs[2] = &mOutputPrefixV; + if (!mActionRunner->wireLanguageOutputs(languageOutputs, stream)) + { + LOG_ERROR("VlaInferenceRuntime(): failed to wire PI0.5 prefix_k/prefix_v into action runner."); + return false; + } + if (!mActionRunner->sampleActions(stream)) + { + LOG_ERROR("VlaInferenceRuntime(): PI0.5 action sampling failed."); + return false; + } + if (!mActionRunner->copyActionsToHost(response.outputActions, stream)) + { + LOG_ERROR("VlaInferenceRuntime(): failed to copy denoised actions to host."); + return false; + } + auto const diffusorEnd = std::chrono::steady_clock::now(); + double const diffusorMs = std::chrono::duration(diffusorEnd - diffusorStart).count(); + LOG_INFO("Stage timings - Diffusor: %.3f ms", diffusorMs); + } + else if (runContextTensorAction) + { + auto const diffusorStart = std::chrono::steady_clock::now(); + if (!mActionRunner->wireStaticInputs(request, stream)) + { + LOG_ERROR("VlaInferenceRuntime(): failed to wire static action inputs (state/embodiment)."); + return false; + } + if (mActionContextRunner) + { + int32_t const actionContextSeqLen = mActionContextRunner->getMaxSeqLen(); + if (!mActionContextRunner->reshapeForContext(activeBatchSize, actionContextSeqLen)) + { + LOG_ERROR("VlaInferenceRuntime(): failed to reshape action context runner."); + return false; + } + if (!mActionContextRunner->copyLmHiddenFrom(mOutputContextEmbeds, stream, actualContextLength)) + { + LOG_ERROR("VlaInferenceRuntime(): failed to copy lm_hidden_states into action context runner."); + return false; + } + if (!mActionContextRunner->resetExecutionContext(stream)) + { + LOG_ERROR("VlaInferenceRuntime(): failed to reset action context execution context."); + return false; + } + if (!mActionContextRunner->infer(stream)) + { + LOG_ERROR("VlaInferenceRuntime(): action context inference failed."); + return false; + } + if (!mActionRunner->copyInputFrom("context_embs", mActionContextRunner->getVlEmbs(), stream)) + { + LOG_ERROR("VlaInferenceRuntime(): failed to wire vl_embs into action runner."); + return false; + } + } + else + { + std::vector languageOutputs; + languageOutputs.push_back(&mOutputLogits); + languageOutputs.push_back(&mOutputContextEmbeds); + if (!mActionRunner->wireLanguageOutputs(languageOutputs, stream)) + { + LOG_ERROR("VlaInferenceRuntime(): failed to wire language context_embs into action runner."); + return false; + } + } + if (!mActionRunner->resetExecutionContext(stream)) + { + LOG_ERROR("VlaInferenceRuntime(): failed to reset action runner execution context."); + return false; + } + if (!mActionRunner->sampleActions(stream)) + { + LOG_ERROR("VlaInferenceRuntime(): context-tensor action sampling failed."); + return false; + } + if (!mActionRunner->copyActionsToHost(response.outputActions, stream)) + { + LOG_ERROR("VlaInferenceRuntime(): failed to copy denoised actions to host."); + return false; + } + auto const diffusorEnd = std::chrono::steady_clock::now(); + double const diffusorMs = std::chrono::duration(diffusorEnd - diffusorStart).count(); + LOG_INFO("Stage timings - Diffusor: %.3f ms", diffusorMs); + } + + auto const e2eEnd = std::chrono::steady_clock::now(); + double const e2eMs = std::chrono::duration(e2eEnd - e2eStart).count(); + LOG_INFO("Stage timings - E2E: %.3f ms", e2eMs); + + return true; +} + +bool VlaInferenceRuntime::captureDecodingCUDAGraph(cudaStream_t stream) +{ + int32_t const maxSupportedBatchSize = mEngineConfig.maxSupportedBatchSize; + int32_t const minSupportedBatchSize = 1; + + bool captureStatus{true}; + // Capture the CUDA graph for all available batch sizes. + for (int32_t batchSize = minSupportedBatchSize; batchSize <= maxSupportedBatchSize; ++batchSize) + { + check::check(mSelectedIndices.reshape({batchSize, 1}), "Tensor reshape failed"); + check::check(mInputsEmbeds.reshape({batchSize, 1, mEngineConfig.hiddenSize}), "Tensor reshape failed"); + check::check(mOutputLogits.reshape({batchSize, mEngineConfig.outputVocabSize}), "Tensor reshape failed"); + + captureStatus &= mLLMEngineRunner->captureVanillaDecodingCudaGraph( + mInputsEmbeds, mOutputLogits, mEmptyLoraWeightsName, stream); + if (mEngineConfig.maxSupportedLoraRank > 0) + { + for (auto const& loraWeightsName : mLLMEngineRunner->getAvailableLoraWeights()) + { + captureStatus &= mLLMEngineRunner->captureVanillaDecodingCudaGraph( + mInputsEmbeds, mOutputLogits, loraWeightsName, stream); + } + } + } + + if (captureStatus) + { + LOG_INFO( + "VlaInferenceRuntime(): Successfully captured the decoding CUDA graph for all execution batch sizes and " + "LoRA weights."); + } + else + { + LOG_WARNING( + "VlaInferenceRuntime(): Failed to capture the decoding CUDA graph for some of execution batch sizes and " + "LoRA weights."); + } + return captureStatus; +} + +VlaInferenceRuntime::TokenCountInfo VlaInferenceRuntime::calculateTokenCounts( + std::vector> const& batchedInputIds, std::vector const& systemPrompts, + std::string const& loraWeightsName) const noexcept +{ + TokenCountInfo tokenCount; + int32_t const activeBatchSize = static_cast(batchedInputIds.size()); + + for (int32_t i = 0; i < activeBatchSize; ++i) + { + int32_t contextLength = static_cast(batchedInputIds[i].size()); + // Calculate reused length from system prompt cache + auto const promptKey = keySystemPromptWithLoraWeights(systemPrompts[i], loraWeightsName); + if (mSystemPromptKVCache.find(promptKey) != mSystemPromptKVCache.end()) + { + int32_t reusedLength = static_cast(mSystemPromptKVCache.at(promptKey).tokenizedPrompt.size()); + tokenCount.totalReusedTokens += reusedLength; + tokenCount.totalComputedTokens += (contextLength - reusedLength); + } + else + { + tokenCount.totalComputedTokens += contextLength; + } + } + + return tokenCount; +} + +bool VlaInferenceRuntime::genAndSaveSystemPromptKVCache( + std::string const& prompt, std::string const& loraWeightsName, cudaStream_t stream) +{ + if (prompt.empty()) + { + LOG_DEBUG("VlaInferenceRuntime(): The prompt is empty. Skip saving system prompt KVCache."); + return true; + } + + // hash the prompt if check if the prompt cache already exists. + auto const promptKey = keySystemPromptWithLoraWeights(prompt, loraWeightsName); + if (mSystemPromptKVCache.find(promptKey) != mSystemPromptKVCache.end()) + { + LOG_DEBUG( + "VlaInferenceRuntime(): The system prompt KVCache already exists for the prompt: {%s}", prompt.c_str()); + return true; + } + + auto tokenizedPrompt = mTokenizer->encode(prompt, true); + if (tokenizedPrompt.empty()) + { + LOG_ERROR("Failed to encode system prompt for KVCache generation."); + return false; + } + int32_t const promptIdsLength = static_cast(tokenizedPrompt.size()); + int32_t const activeBatchSize = 1; + + if (promptIdsLength > mEngineConfig.maxSupportedInputLength) + { + LOG_ERROR( + "VlaInferenceRuntime(): The prompt length (%d) exceeds the max supported input length (%d) of the LLM " + "Engine.", + promptIdsLength, mEngineConfig.maxSupportedInputLength); + return false; + } + + std::vector> batchedInputIds(activeBatchSize, tokenizedPrompt); + std::vector batchedSystemPrompts(activeBatchSize, prompt); + if (!setUpForPrefillExecution(batchedInputIds, batchedSystemPrompts, loraWeightsName, stream)) + { + LOG_ERROR( + "VlaInferenceRuntime(): Prefill execution setup failed. Cannot generate the KVCache for this prompt."); + return false; + } + + // Execute prefill step to initialize the KVCache data. + // Perform embedding lookup + int32_t const prefillSequenceLength = mInputIds.getShape()[1]; + check::check(mInputsEmbeds.reshape({activeBatchSize, prefillSequenceLength, mEngineConfig.hiddenSize}), + "Tensor reshape failed"); + + // Get embeddings from independent runners + rt::OptionalInputTensor visionEmbeddings + = mVisionRunner ? std::optional{std::ref(mVisionRunner->getOutputEmbedding())} : std::nullopt; + rt::OptionalInputTensor audioEmbeddings + = mAudioRunner ? std::optional{std::ref(mAudioRunner->getOutputEmbedding())} : std::nullopt; + + if (audioEmbeddings.has_value()) + { + // Audio present: use embeddingLookupMultimodal (handles audio and/or vision) + auto const inputShape = mInputIds.getShape(); + size_t const inputSizeBytes = inputShape.volume() * sizeof(int32_t); + rt::Tensor inputIdsCPU(inputShape, rt::DeviceType::kCPU, mInputIds.getDataType()); + CUDA_CHECK( + cudaMemcpy(inputIdsCPU.rawPointer(), mInputIds.rawPointer(), inputSizeBytes, cudaMemcpyDeviceToHost)); + + std::optional audioTokenId + = (mEngineConfig.audioTokenId != 0) ? std::optional{mEngineConfig.audioTokenId} : std::nullopt; + std::optional imageTokenId + = (mEngineConfig.imageTokenId != 0) ? std::optional{mEngineConfig.imageTokenId} : std::nullopt; + rt::Tensor multimodalIndicesCPU + = generateMultimodalIndices(inputIdsCPU, audioTokenId, imageTokenId, mEngineConfig.vocabSize); + + auto const indicesShape = multimodalIndicesCPU.getShape(); + size_t const indicesSizeBytes = indicesShape.volume() * sizeof(int32_t); + mMultimodalIndices = rt::Tensor(indicesShape, rt::DeviceType::kGPU, multimodalIndicesCPU.getDataType()); + CUDA_CHECK(cudaMemcpy(mMultimodalIndices.rawPointer(), multimodalIndicesCPU.rawPointer(), indicesSizeBytes, + cudaMemcpyHostToDevice)); + + kernel::embeddingLookupMultimodal(mInputIds, mEmbedding.table, mEmbedding.scalesAsOptional(), + std::optional{std::ref(mMultimodalIndices)}, imageTokenId, visionEmbeddings, audioTokenId, audioEmbeddings, + mInputsEmbeds, stream); + } + else if (visionEmbeddings.has_value()) + { + // Vision-only (Qwen2-VL, InternVL, etc.) + rt::Tensor const& imageEmbedsTensor = visionEmbeddings.value().get(); + kernel::embeddingLookupWithImageInsertion( + mInputIds, mEmbedding.table, mEmbedding.scalesAsOptional(), imageEmbedsTensor, mInputsEmbeds, stream); + } + else + { + // Standard embedding lookup (pure text) + kernel::embeddingLookup(mInputIds, mEmbedding.table, mEmbedding.scalesAsOptional(), mInputsEmbeds, stream); + } + + // Process deepstack features: perform embedding lookup if vision runner is available + rt::OptionalInputTensors deepstackEmbeds{}; + if (mEngineConfig.numDeepstackFeatures > 0 && mVisionRunner) + { + rt::OptionalInputTensors deepstackFeatures = mVisionRunner->getDeepstackFeatures(); + + rt::OptionalInputTensor deepstackMultimodalIndices{std::nullopt}; + if (mMultimodalIndices.getShape().volume() > 0) + { + deepstackMultimodalIndices = std::ref(mMultimodalIndices); + } + + for (int32_t idx = 0; idx < static_cast(deepstackFeatures.size()); ++idx) + { + rt::Tensor const& featureTensor = deepstackFeatures[idx].get(); + + check::check( + mDeepstackEmbeds[idx].reshape({activeBatchSize, prefillSequenceLength, mEngineConfig.hiddenSize}), + "Tensor reshape failed"); + kernel::assembleDeepstackEmbedding(mInputIds, featureTensor, mEngineConfig.vocabSize, mDeepstackEmbeds[idx], + stream, mEngineConfig.imageTokenId, deepstackMultimodalIndices); + + deepstackEmbeds.push_back(std::ref(mDeepstackEmbeds[idx])); + } + } + + rt::OptionalOutputTensor outputHiddenStates{std::nullopt}; + rt::OptionalOutputTensor outputContextEmbeds{std::nullopt}; + if (mEngineConfig.enableContextEmb || mEngineConfig.enableLmHiddenStates) + { + outputContextEmbeds = std::ref(mOutputContextEmbeds); + } + bool prefillStatus = mLLMEngineRunner->executePrefillStep(mInputsEmbeds, mHostContextLengths, deepstackEmbeds, + mOutputLogits, outputHiddenStates, stream, outputContextEmbeds); + if (!prefillStatus) + { + LOG_ERROR("VlaInferenceRuntime(): Failed to execute prefill step."); + return false; + } + + // Copy out the KVCache content from the prefill step. + auto& linearKVCache = mLLMEngineRunner->getLinearKVCache(); + auto cacheConfig = linearKVCache.getConfig(); + auto kvCacheBuffer = linearKVCache.getKVCacheBuffer(); + rt::Coords savedKVCacheShape{ + cacheConfig.numAttentionLayers, 2, cacheConfig.numKVHeads, promptIdsLength, cacheConfig.headDim}; + + VlaSystemPromptKVCache savedKVCache; + savedKVCache.systemPrompt = prompt; + savedKVCache.tokenizedPrompt = tokenizedPrompt; + savedKVCache.kvCacheContent = rt::Tensor(savedKVCacheShape, rt::DeviceType::kGPU, + linearKVCache.getConfig().kvCacheTypeTRT, "VlaInferenceRuntime::savedKVCache.kvCacheContent"); + + // We only process one sequence at a time. + constexpr int32_t CACHE_BATCH_IDX{0}; + kernel::saveKVCacheIntoTensor(savedKVCache.kvCacheContent, kvCacheBuffer, CACHE_BATCH_IDX, stream); + + // Save recurrent and conv states for hybrid layers + if (mEngineConfig.numLinearAttnLayers > 0) + { + savedKVCache.recurrentStateContents + = mLLMEngineRunner->getLinearKVCache().captureRecurrentStates(CACHE_BATCH_IDX, stream); + savedKVCache.convStateContents + = mLLMEngineRunner->getLinearKVCache().captureConvStates(CACHE_BATCH_IDX, stream); + } + + mSystemPromptKVCache.insert({promptKey, std::move(savedKVCache)}); + + CUDA_CHECK(cudaStreamSynchronize(stream)); + LOG_DEBUG("VlaInferenceRuntime(): The KVCache is saved for the prompt: {%s}", prompt.c_str()); + + return true; +} + +} // namespace rt +} // namespace trt_edgellm diff --git a/cpp/runtime/vlaInferenceRuntime.h b/cpp/runtime/vlaInferenceRuntime.h new file mode 100644 index 00000000..6b82c9ab --- /dev/null +++ b/cpp/runtime/vlaInferenceRuntime.h @@ -0,0 +1,198 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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. + */ + +#pragma once + +#include "action/actionContextRunner.h" +#include "action/actionRunner.h" +#include "common/hashUtils.h" +#include "multimodal/multimodalRunner.h" +#include "profiling/metrics.h" +#include "profiling/timer.h" +#include "runtime/llmEngineRunner.h" +#include "runtime/llmRuntimeUtils.h" +#include "tokenizer/tokenizer.h" +#include +#include +#include + +namespace trt_edgellm +{ +namespace rt +{ + +/*! \brief Structure to hold cached system prompt and its KV cache + */ +struct VlaSystemPromptKVCache +{ + std::string systemPrompt; //!< The system prompt text + std::vector tokenizedPrompt; //!< Tokenized version of the system prompt + rt::Tensor kvCacheContent; //!< Cached KV cache content for the system prompt + std::vector recurrentStateContents; //!< Cached recurrent states for hybrid layers + std::vector convStateContents; //!< Cached conv states for hybrid layers +}; + +/*! \brief LLM Inference Runtime for handling generation requests + */ +class VlaInferenceRuntime +{ +public: + /*! \brief Construct an LLM Inference Runtime + * \param engineDir Directory containing the LLM engine + * \param multimodalEngineDir Directory containing the multimodal engine + * \param loraWeightsMap Map of LoRA weights names to their paths + * \param stream CUDA stream for initialization + * \throws std::runtime_error if loading data from directories fails, or tensor shapes are invalid, or there is an + * initialization error + */ + VlaInferenceRuntime(std::string const& engineDir, std::string const& multimodalEngineDir, + std::unordered_map const& loraWeightsMap, cudaStream_t stream); + + /*! \brief Destructor + */ + ~VlaInferenceRuntime() noexcept = default; + + /*! \brief Handle an LLM generation request + * \param request The generation request containing prompt and generation parameters + * \param response The generation response to be filled with output + * \param stream CUDA stream for execution + * \return True if request was handled successfully, false otherwise + * \throws std::runtime_error if a LLM operation or CUDA operation fails + */ + bool handleRequest(LLMGenerationRequest const& request, LLMGenerationResponse& response, cudaStream_t stream); + + /*! \brief Capture CUDA graph for the decoding step to optimize performance + * \param stream CUDA stream for graph capture + * \return True if graph was captured successfully, false otherwise + * \throws std::runtime_error if a CUDA operation fails + */ + bool captureDecodingCUDAGraph(cudaStream_t stream); + + /*! \brief Execute the prefill step generation of the KVCache for the prompt and save for later usage + * + * \param prompt The system prompt to generate the KVCache + * \param loraWeightsName The name of the LoRA weights + * \param stream The CUDA stream used for the generation + * \return True if the KVCache is generated and saved successfully, false otherwise + * \throws std::runtime_error if a CUDA operation fails + */ + bool genAndSaveSystemPromptKVCache( + std::string const& prompt, std::string const& loraWeightsName, cudaStream_t stream); + + /*! \brief Get LLM prefill stage metrics + * \return Reference to prefill metrics + */ + metrics::LLMPrefillMetrics const& getPrefillMetrics() const noexcept + { + return mPrefillMetrics; + } + + /*! \brief Get LLM generation stage metrics + * \return Reference to generation metrics + */ + metrics::LLMGenerationMetrics const& getGenerationMetrics() const noexcept + { + return mGenerationMetrics; + } + + /*! \brief Set the random seed used when initializing the action diffusion noise trajectory + * \param seed Random seed value; has no effect if no action runner is loaded + */ + void setActionNoiseSeed(int32_t seed) noexcept; + + /*! \brief Get multimodal metrics (returns empty metrics if no multimodal runner) + * \return Multimodal metrics, or empty metrics if no multimodal runner is available + */ + metrics::MultimodalMetrics getMultimodalMetrics() const noexcept + { + return mVisionRunner ? mVisionRunner->getMultimodalMetrics() + : mAudioRunner ? mAudioRunner->getMultimodalMetrics() + : metrics::MultimodalMetrics{}; + } + +private: + /*! \brief Helper structure to hold token counting results + */ + struct TokenCountInfo + { + int32_t totalReusedTokens{0}; //!< Number of tokens reused from KV cache + int32_t totalComputedTokens{0}; //!< Number of tokens that need computation + }; + + //! Calculate token counts (reused vs computed) for performance tracking. + //! \param batchedInputIds Batched input token IDs + //! \param systemPrompts System prompts for each batch element + //! \param loraWeightsName Name of the LoRA weights being used + //! \return TokenCountInfo structure containing reused and computed token counts + TokenCountInfo calculateTokenCounts(std::vector> const& batchedInputIds, + std::vector const& systemPrompts, std::string const& loraWeightsName) const noexcept; + + rt::Tensor mSharedExecContextMemory{}; //!< Shared device memory for LLM and multimodal execution contexts + std::unique_ptr mLLMEngineRunner{nullptr}; //!< LLM engine runner instance + std::unique_ptr mAudioRunner{nullptr}; //!< Audio runner instance (optional) + std::unique_ptr mVisionRunner{nullptr}; //!< Vision runner instance (optional) + std::unique_ptr mActionRunner{ + nullptr}; //!< Action/diffusion head runner (optional; prefix-KV or context-tensor handoff) + std::unique_ptr mActionContextRunner{ + nullptr}; //!< GR00T action-context projection engine (lm_hidden_states -> vl_embs) + std::unique_ptr mTokenizer{nullptr}; //!< Tokenizer instance + hash_utils::HashMap, VlaSystemPromptKVCache> + mSystemPromptKVCache{}; //!< Cache of system prompts / LORA weights and their KV caches + + EmbeddingData mEmbedding{}; //!< Shared embedding table [vocabSize, hiddenSize] and optional FP8 scales + rt::Tensor mSamplingWorkspace{}; //!< Workspace tensor for sampling operations + rt::Tensor mInputIds{}; //!< Input token IDs tensor + rt::Tensor mInputsEmbeds{}; //!< Input embeddings tensor [batchSize, seqLen, hiddenSize] + rt::Tensor mMultimodalIndices{}; //!< Multimodal indices tensor [batchSize, seqLen] for audio/image embeddings + std::vector mDeepstackEmbeds; //!< Deepstack embeddings tensors for Qwen3-VL (one per feature) + rt::Tensor mHostPackedInputIds{}; //!< Host tensor for packed input IDs + rt::Tensor mHostContextLengths{}; //!< Host tensor for context lengths + rt::Tensor mOutputLogits{}; //!< Output logits tensor + rt::Tensor mSelectedIndices{}; //!< Selected token indices tensor + rt::Tensor mHostSelectedTokenIds{}; //!< Host tensor for selected token IDs + rt::Tensor mHostReuseKVCacheLengths{}; //!< Reuse KV cache lengths for prefill + rt::Tensor mVocabMappingTable{}; //!< Vocab mapping table for reduced vocab (empty if not used) + rt::Tensor mOutputContextEmbeds{}; //!< Language prefill sequence output (context_embs or lm_hidden_states) + rt::Tensor mOutputPrefixK{}; //!< PI0.5 stacked prefix K cache for action head + rt::Tensor mOutputPrefixV{}; //!< PI0.5 stacked prefix V cache for action head + std::string mEmptyLoraWeightsName{""}; //!< Empty LoRA weights name for default case + + bool mUseCompactPrefixPadding{false}; //!< Pad compact VLA prefixes to engine max_seq_len + + LLMEngineRunnerConfig mEngineConfig{}; //!< Engine configuration + + metrics::LLMPrefillMetrics mPrefillMetrics; //!< Stage-specific metrics to store number of tokens in prefill + metrics::LLMGenerationMetrics + mGenerationMetrics; //!< Stage-specific metrics to store number of tokens in generation + + //! Examine and validate the generation request. + //! \param request The generation request to examine + //! \return True if request is valid, false otherwise + bool examineRequest(LLMGenerationRequest const& request) noexcept; + + //! Set up tensors and state for prefill execution. + //! \param batchedInputIds Batched input token IDs + //! \param systemPrompts System prompts for each batch element + //! \param loraWeightsName Name of the LoRA weights being used + //! \param stream CUDA stream for execution + //! \return True if setup was successful, false otherwise + //! \throws std::runtime_error if system prompt is malformed, or a CUDA operation fails + bool setUpForPrefillExecution(std::vector> const& batchedInputIds, + std::vector const& systemPrompts, std::string const& loraWeightsName, cudaStream_t stream); +}; +} // namespace rt +} // namespace trt_edgellm diff --git a/cpp/runtime/wfmInferenceRuntime.cpp b/cpp/runtime/wfmInferenceRuntime.cpp new file mode 100644 index 00000000..baa41ea7 --- /dev/null +++ b/cpp/runtime/wfmInferenceRuntime.cpp @@ -0,0 +1,426 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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. + */ + +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/wfmInferenceRuntime.h" + +#include "common/checkMacros.h" +#include "common/cudaUtils.h" +#include "common/logger.h" +#include "runtime/llmRuntimeUtils.h" +#include "tokenizer/tokenizer.h" + +#include +#include +#include + +namespace trt_edgellm +{ +namespace rt +{ +namespace +{ + +bool denoiseComponentsPresent(std::filesystem::path const& root) +{ + return std::filesystem::exists(root / "embed" / "config.json") + && std::filesystem::exists(root / "mot_backbone" / "config.json") + && std::filesystem::exists(root / "denoise_head" / "config.json"); +} + +bool componentPresent(std::filesystem::path const& root, char const* name) +{ + return std::filesystem::exists(root / name / "config.json"); +} + +std::optional soundLatentShapeFromPacked(CosmosEngineConfig const& config, CosmosPackedStatic const& packed) +{ + if (packed.soundTokenShape.empty()) + { + return std::nullopt; + } + + if (packed.soundTokenShape.size() == 1) + { + return rt::Coords({1, config.soundDim, packed.soundTokenShape[0]}); + } + if (packed.soundTokenShape.size() == 2) + { + return rt::Coords({1, packed.soundTokenShape[0], packed.soundTokenShape[1]}); + } + + return std::nullopt; +} + +} // namespace + +WFMInferenceRuntime::WFMInferenceRuntime(std::string const& engineDir, cudaStream_t stream) + : mEngineDir(engineDir) +{ + std::filesystem::path const root(engineDir); + mConfig = loadCosmosEngineConfig(root / "config.json"); + mPackedStatic = loadCosmosPackedStatic(root / "packing_static.json"); + + mTokenizer = std::make_unique(); + std::filesystem::path tokenizerDir = root / "tokenizer"; + if (!std::filesystem::exists(tokenizerDir / "tokenizer.json")) + { + tokenizerDir = root; + } + if (!mTokenizer->loadFromHF(tokenizerDir)) + { + throw std::runtime_error("WFMInferenceRuntime: failed to load tokenizer from " + tokenizerDir.string()); + } + + mEmbedding = loadEmbeddingTable(root / "embedding.safetensors", stream); + + int32_t const maxUndLen = mPackedStatic.sequenceLength; + int32_t const maxGenLen + = mPackedStatic.visionTokenShape[0] * mPackedStatic.visionTokenShape[1] * mPackedStatic.visionTokenShape[2]; + mTextPhase0.undSeq = rt::Tensor({maxUndLen, mConfig.hiddenSize}, rt::DeviceType::kGPU, nvinfer1::DataType::kHALF, + "WFMInferenceRuntime::undSeq"); + mTextPhase0.cosUnd = rt::Tensor( + {maxUndLen, mConfig.headDim}, rt::DeviceType::kGPU, nvinfer1::DataType::kHALF, "WFMInferenceRuntime::cosUnd"); + mTextPhase0.sinUnd = rt::Tensor( + {maxUndLen, mConfig.headDim}, rt::DeviceType::kGPU, nvinfer1::DataType::kHALF, "WFMInferenceRuntime::sinUnd"); + mTextPhase0.cosGen = rt::Tensor( + {maxGenLen, mConfig.headDim}, rt::DeviceType::kGPU, nvinfer1::DataType::kHALF, "WFMInferenceRuntime::cosGen"); + mTextPhase0.sinGen = rt::Tensor( + {maxGenLen, mConfig.headDim}, rt::DeviceType::kGPU, nvinfer1::DataType::kHALF, "WFMInferenceRuntime::sinGen"); + + std::string const visualEncodeDir = (root / "visual_encode").string(); + std::string const visualDecodeDir = (root / "visual_decode").string(); + mVaeEncodeRunner = std::make_unique(visualEncodeDir, stream); + mVaeDecodeRunner = std::make_unique(visualDecodeDir, stream); + + if (componentPresent(root, "audio_encode")) + { + mAudioEncodeRunner = std::make_unique((root / "audio_encode").string(), stream); + } + else if (mConfig.enableSound) + { + LOG_WARNING("WFMInferenceRuntime: enable_sound=true but audio_encode/ is missing."); + } + + if (componentPresent(root, "audio_decode")) + { + mAudioDecodeRunner = std::make_unique((root / "audio_decode").string(), stream); + } + else if (mConfig.enableSound) + { + LOG_WARNING("WFMInferenceRuntime: enable_sound=true but audio_decode/ is missing."); + } + + if (denoiseComponentsPresent(root)) + { + mDenoiseRunner = std::make_unique(root, mConfig, stream); + } + else + { + LOG_WARNING("WFMInferenceRuntime: denoise engines incomplete under %s; denoise path will be unavailable.", + engineDir.c_str()); + } + + int64_t sharedContextMemorySize + = std::max(mVaeEncodeRunner->getRequiredContextMemorySize(), mVaeDecodeRunner->getRequiredContextMemorySize()); + if (mAudioEncodeRunner) + { + sharedContextMemorySize = std::max(sharedContextMemorySize, mAudioEncodeRunner->getRequiredContextMemorySize()); + } + if (mAudioDecodeRunner) + { + sharedContextMemorySize = std::max(sharedContextMemorySize, mAudioDecodeRunner->getRequiredContextMemorySize()); + } + if (mDenoiseRunner) + { + sharedContextMemorySize = std::max(sharedContextMemorySize, mDenoiseRunner->getRequiredContextMemorySize()); + } + mSharedExecContextMemory = rt::Tensor({sharedContextMemorySize}, rt::DeviceType::kGPU, nvinfer1::DataType::kUINT8, + "WFMInferenceRuntime::mSharedExecContextMemory"); + if (!mVaeEncodeRunner->setContextMemory(mSharedExecContextMemory)) + { + throw std::runtime_error("WFMInferenceRuntime: failed to set shared context memory for visual_encode"); + } + if (!mVaeDecodeRunner->setContextMemory(mSharedExecContextMemory)) + { + throw std::runtime_error("WFMInferenceRuntime: failed to set shared context memory for visual_decode"); + } + if (mAudioEncodeRunner && !mAudioEncodeRunner->setContextMemory(mSharedExecContextMemory)) + { + throw std::runtime_error("WFMInferenceRuntime: failed to set shared context memory for audio_encode"); + } + if (mAudioDecodeRunner && !mAudioDecodeRunner->setContextMemory(mSharedExecContextMemory)) + { + throw std::runtime_error("WFMInferenceRuntime: failed to set shared context memory for audio_decode"); + } + if (mDenoiseRunner && !mDenoiseRunner->setContextMemory(mSharedExecContextMemory)) + { + throw std::runtime_error("WFMInferenceRuntime: failed to set shared context memory for denoise"); + } + + mOutputVideo = std::make_shared(mVaeDecodeRunner->getPixelShape(), rt::DeviceType::kGPU, + mVaeDecodeRunner->getPixels().getDataType(), "WFMInferenceRuntime::mOutputVideo"); + + mNoisyLatents = rt::Tensor(mVaeEncodeRunner->getLatentShape(), rt::DeviceType::kGPU, + mVaeEncodeRunner->getLatents().getDataType(), "WFMInferenceRuntime::mNoisyLatents"); + + if (mAudioDecodeRunner) + { + mOutputWaveform = std::make_shared(mAudioDecodeRunner->getWaveformShape(), rt::DeviceType::kGPU, + mAudioDecodeRunner->getWaveform().getDataType(), "WFMInferenceRuntime::mOutputWaveform"); + } + + rt::Coords soundLatentShape{}; + nvinfer1::DataType soundLatentDtype{nvinfer1::DataType::kHALF}; + if (mAudioEncodeRunner) + { + soundLatentShape = mAudioEncodeRunner->getSoundLatentShape(); + soundLatentDtype = mAudioEncodeRunner->getSoundLatents().getDataType(); + } + else if (auto packedShape = soundLatentShapeFromPacked(mConfig, mPackedStatic)) + { + soundLatentShape = *packedShape; + } + else if (mAudioDecodeRunner) + { + soundLatentShape = mAudioDecodeRunner->getSoundLatentShape(); + soundLatentDtype = mAudioDecodeRunner->getSoundLatents().getDataType(); + } + + if (soundLatentShape.getNumDims() == 3) + { + mNoisySoundLatents = rt::Tensor( + soundLatentShape, rt::DeviceType::kGPU, soundLatentDtype, "WFMInferenceRuntime::mNoisySoundLatents"); + mZeroSoundLatents = rt::Tensor( + soundLatentShape, rt::DeviceType::kGPU, soundLatentDtype, "WFMInferenceRuntime::mZeroSoundLatents"); + CUDA_CHECK(cudaMemsetAsync(mZeroSoundLatents.rawPointer(), 0, + static_cast(mZeroSoundLatents.getShape().volume()) + * rt::utils::getTypeSize(mZeroSoundLatents.getDataType()), + stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + } + + LOG_INFO( + "WFMInferenceRuntime loaded from %s (frames=%d, %dx%d, hidden=%d, seq_len=%d, denoise=%s, sound=%s, " + "shared_ctx=%zu bytes)", + engineDir.c_str(), mConfig.numFrames, mConfig.height, mConfig.width, mConfig.hiddenSize, + mPackedStatic.sequenceLength, mDenoiseRunner ? "yes" : "no", mConfig.enableSound ? "enabled" : "disabled", + static_cast(sharedContextMemorySize)); +} + +bool WFMInferenceRuntime::examineRequest(WFMGenerationRequest const& request) const noexcept +{ + return examineWFMRequest(request, mConfig, mPackedStatic); +} + +bool WFMInferenceRuntime::prepareTextPhase0(std::string const& prompt, cudaStream_t stream) +{ + return prepareCosmosTextPhase0(*mTokenizer, mEmbedding, mConfig, mPackedStatic, prompt, mTextPhase0, stream); +} + +bool WFMInferenceRuntime::allocateZeroSoundLatents( + rt::Coords const& shape, nvinfer1::DataType dtype, cudaStream_t stream) +{ + mZeroSoundLatents = rt::Tensor(shape, rt::DeviceType::kGPU, dtype, "WFMInferenceRuntime::mZeroSoundLatents"); + mNoisySoundLatents = rt::Tensor(shape, rt::DeviceType::kGPU, dtype, "WFMInferenceRuntime::mNoisySoundLatents"); + CUDA_CHECK(cudaMemsetAsync(mZeroSoundLatents.rawPointer(), 0, + static_cast(shape.volume()) * rt::utils::getTypeSize(dtype), stream)); + return true; +} + +bool WFMInferenceRuntime::handleRequest( + WFMGenerationRequest const& request, WFMGenerationResponse& response, cudaStream_t stream) +{ + response = WFMGenerationResponse{}; + + if (!examineRequest(request)) + { + LOG_ERROR("WFMInferenceRuntime: request validation failed."); + return false; + } + + if (!prepareTextPhase0(request.prompt, stream)) + { + LOG_ERROR("WFMInferenceRuntime: Phase 0 text preparation failed."); + return false; + } + + if (!mVaeEncodeRunner->copyPixelsFrom(request.pixels, stream)) + { + LOG_ERROR("WFMInferenceRuntime: failed to copy pixels into visual_encode runner."); + return false; + } + + if (!mVaeEncodeRunner->encode(stream)) + { + LOG_ERROR("WFMInferenceRuntime: visual_encode TRT inference failed."); + return false; + } + + if (!seedNoisyVisionLatents(mVaeEncodeRunner->getLatents(), mNoisyLatents, mPackedStatic.visionNoisyFrameIndexes, + request.seed, 1.F, stream)) + { + LOG_ERROR("WFMInferenceRuntime: failed to seed noisy vision latents."); + return false; + } + + rt::Tensor const* cleanSoundLatents = nullptr; + if (request.inputWaveform.buffer) + { + if (!mAudioEncodeRunner) + { + LOG_ERROR("WFMInferenceRuntime: input waveform provided but audio_encode engine is not loaded."); + return false; + } + if (!mAudioEncodeRunner->copyWaveformFrom(request.inputWaveform, stream)) + { + LOG_ERROR("WFMInferenceRuntime: failed to copy waveform into audio_encode runner."); + return false; + } + if (!mAudioEncodeRunner->encode(stream)) + { + LOG_ERROR("WFMInferenceRuntime: audio_encode TRT inference failed."); + return false; + } + cleanSoundLatents = &mAudioEncodeRunner->getSoundLatents(); + if (mNoisySoundLatents.getShape().volume() == 0) + { + if (!allocateZeroSoundLatents(cleanSoundLatents->getShape(), cleanSoundLatents->getDataType(), stream)) + { + LOG_ERROR("WFMInferenceRuntime: failed to allocate sound latent buffers."); + return false; + } + } + } + + bool const runSoundDenoise = request.generateSound && mConfig.enableSound; + if (runSoundDenoise) + { + if (!mDenoiseRunner || !mDenoiseRunner->hasSoundPath()) + { + LOG_ERROR("WFMInferenceRuntime: generateSound requested but denoise_head_sound is unavailable."); + return false; + } + if (mNoisySoundLatents.getShape().volume() == 0) + { + LOG_ERROR("WFMInferenceRuntime: sound latent buffers are not allocated."); + return false; + } + + rt::Tensor const& soundSeedSource = cleanSoundLatents != nullptr ? *cleanSoundLatents : mZeroSoundLatents; + if (!seedNoisySoundLatents( + soundSeedSource, mNoisySoundLatents, mPackedStatic.soundNoisySlotIndexes, request.seed, 1.F, stream)) + { + LOG_ERROR("WFMInferenceRuntime: failed to seed noisy sound latents."); + return false; + } + } + + if (mDenoiseRunner) + { + int32_t const numSteps = request.numInferenceSteps > 0 ? request.numInferenceSteps : mConfig.numInferenceSteps; + + CosmosDenoiseLatents denoiseLatents{}; + denoiseLatents.vision = &mNoisyLatents; + if (runSoundDenoise) + { + denoiseLatents.sound = &mNoisySoundLatents; + } + + if (!mDenoiseRunner->sampleLatents(mTextPhase0, denoiseLatents, numSteps, stream)) + { + LOG_ERROR("WFMInferenceRuntime: denoise failed."); + return false; + } + } + else + { + LOG_WARNING( + "WFMInferenceRuntime: CosmosDenoiseRunner unavailable; skipping denoise (encode->seed->decode smoke " + "path)."); + } + + if (!mVaeDecodeRunner->copyLatentsFrom(mNoisyLatents, stream)) + { + LOG_ERROR("WFMInferenceRuntime: failed to copy latents into visual_decode runner."); + return false; + } + + if (!mVaeDecodeRunner->decode(stream)) + { + LOG_ERROR("WFMInferenceRuntime: visual_decode TRT inference failed."); + return false; + } + + if (!mVaeDecodeRunner->copyPixelsTo(*mOutputVideo, stream)) + { + LOG_ERROR("WFMInferenceRuntime: failed to copy decoded pixels into output buffer."); + return false; + } + + response.outputVideo.buffer = mOutputVideo; + response.outputVideo.batch = 1; + response.outputVideo.channels = 3; + response.outputVideo.numFrames = mConfig.numFrames; + response.outputVideo.height = mConfig.height; + response.outputVideo.width = mConfig.width; + + if (runSoundDenoise && mAudioDecodeRunner && mOutputWaveform) + { + if (!mAudioDecodeRunner->copySoundLatentsFrom(mNoisySoundLatents, stream)) + { + LOG_ERROR("WFMInferenceRuntime: failed to copy sound latents into audio_decode runner."); + return false; + } + if (!mAudioDecodeRunner->decode(stream)) + { + LOG_ERROR("WFMInferenceRuntime: audio_decode TRT inference failed."); + return false; + } + if (!mAudioDecodeRunner->copyWaveformTo(*mOutputWaveform, stream)) + { + LOG_ERROR("WFMInferenceRuntime: failed to copy decoded waveform into output buffer."); + return false; + } + + response.outputWaveform.buffer = mOutputWaveform; + response.outputWaveform.batch = 1; + response.outputWaveform.sampleRate = mConfig.sampleRate; + auto const& waveformShape = mOutputWaveform->getShape(); + if (waveformShape.getNumDims() == 3) + { + response.outputWaveform.numSamples = waveformShape[2]; + } + else if (waveformShape.getNumDims() == 2) + { + response.outputWaveform.numSamples = waveformShape[1]; + } + } + + LOG_INFO("WFMInferenceRuntime: inference complete (video=%s, sound=%s, denoise=%s).", + mOutputVideo->getShape().formatString().c_str(), + response.outputWaveform.buffer ? mOutputWaveform->getShape().formatString().c_str() : "n/a", + mDenoiseRunner ? "yes" : "no"); + return true; +} + +} // namespace rt +} // namespace trt_edgellm diff --git a/cpp/runtime/wfmInferenceRuntime.h b/cpp/runtime/wfmInferenceRuntime.h new file mode 100644 index 00000000..c24f4e1d --- /dev/null +++ b/cpp/runtime/wfmInferenceRuntime.h @@ -0,0 +1,92 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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. + */ + +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "runtime/audioDecodeRunner.h" +#include "runtime/audioEncodeRunner.h" +#include "runtime/cosmosDenoiseRunner.h" +#include "runtime/vaeDecodeRunner.h" +#include "runtime/vaeEncodeRunner.h" +#include "runtime/wfmRuntimeUtils.h" + +#include +#include +#include + +namespace trt_edgellm +{ +namespace tokenizer +{ +class Tokenizer; +} // namespace tokenizer + +namespace rt +{ + +//! Orchestrates Cosmos TRT engines: encode → denoise → decode (vision + optional sound for Omni). +class WFMInferenceRuntime +{ +public: + //! \p engineDir must contain config.json and packing_static.json. + explicit WFMInferenceRuntime(std::string const& engineDir, cudaStream_t stream); + + ~WFMInferenceRuntime() noexcept = default; + + bool handleRequest(WFMGenerationRequest const& request, WFMGenerationResponse& response, cudaStream_t stream); + + CosmosEngineConfig const& getEngineConfig() const noexcept + { + return mConfig; + } + + CosmosPackedStatic const& getPackedStatic() const noexcept + { + return mPackedStatic; + } + +private: + bool examineRequest(WFMGenerationRequest const& request) const noexcept; + bool prepareTextPhase0(std::string const& prompt, cudaStream_t stream); + bool allocateZeroSoundLatents(rt::Coords const& shape, nvinfer1::DataType dtype, cudaStream_t stream); + + std::string mEngineDir; + CosmosEngineConfig mConfig{}; + CosmosPackedStatic mPackedStatic{}; + std::unique_ptr mTokenizer; + EmbeddingData mEmbedding{}; + CosmosTextPhase0 mTextPhase0{}; + std::unique_ptr mVaeEncodeRunner; + std::unique_ptr mVaeDecodeRunner; + std::unique_ptr mAudioEncodeRunner; + std::unique_ptr mAudioDecodeRunner; + std::unique_ptr mDenoiseRunner; + std::shared_ptr mOutputVideo; + std::shared_ptr mOutputWaveform; + rt::Tensor mNoisyLatents{}; + rt::Tensor mNoisySoundLatents{}; + rt::Tensor mZeroSoundLatents{}; + rt::Tensor mSharedExecContextMemory{}; +}; + +} // namespace rt +} // namespace trt_edgellm diff --git a/cpp/runtime/wfmPackingUtils.cpp b/cpp/runtime/wfmPackingUtils.cpp new file mode 100644 index 00000000..7b8fbe39 --- /dev/null +++ b/cpp/runtime/wfmPackingUtils.cpp @@ -0,0 +1,624 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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. + */ + +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/wfmRuntimeUtils.h" + +#include "common/cudaUtils.h" +#include "common/logger.h" +#include "kernels/embeddingKernels/embeddingKernels.h" +#include "tokenizer/tokenizer.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace trt_edgellm +{ +namespace rt +{ +namespace +{ + +std::pair, int64_t> get3dMropeIdsTextTokens( + int32_t numTokens, int64_t temporalOffset, bool useFloatPositions) +{ + std::vector ids(static_cast(numTokens) * 3U); + for (int32_t i = 0; i < numTokens; ++i) + { + int64_t const value = temporalOffset + i; + for (int32_t axis = 0; axis < 3; ++axis) + { + ids[static_cast(axis) * static_cast(numTokens) + static_cast(i)] + = value; + } + } + (void) useFloatPositions; + return {ids, temporalOffset + numTokens}; +} + +std::pair, int64_t> get3dMropeIdsVaeTokens(int32_t gridT, int32_t gridH, int32_t gridW, + int64_t temporalOffset, bool resetSpatialIndices, float fps, float baseFps, int32_t temporalCompressionFactor, + bool enableFpsModulation) +{ + int32_t const numTokens = gridT * gridH * gridW; + std::vector ids(static_cast(numTokens) * 3U); + bool const fpsModulationEnabled = enableFpsModulation && gridT > 1 && fps > 0.F; + + int32_t tokenIdx = 0; + for (int32_t t = 0; t < gridT; ++t) + { + for (int32_t h = 0; h < gridH; ++h) + { + for (int32_t w = 0; w < gridW; ++w) + { + int64_t tIndex = 0; + if (fpsModulationEnabled) + { + float const tps = fps / static_cast(temporalCompressionFactor); + float const baseTps = baseFps / static_cast(temporalCompressionFactor); + float const scaledT = (static_cast(t) / tps) * baseTps + static_cast(temporalOffset); + tIndex = static_cast(std::floor(scaledT)); + } + else + { + tIndex = static_cast(t) + temporalOffset; + } + + int64_t hIndex = h; + int64_t wIndex = w; + if (!resetSpatialIndices) + { + hIndex += temporalOffset; + wIndex += temporalOffset; + } + + ids[static_cast(tokenIdx)] = tIndex; + ids[static_cast(numTokens) + static_cast(tokenIdx)] = hIndex; + ids[static_cast(numTokens) * 2U + static_cast(tokenIdx)] = wIndex; + ++tokenIdx; + } + } + } + + int64_t maxPosition = 0; + for (auto const value : ids) + { + maxPosition = std::max(maxPosition, value); + } + return {ids, maxPosition + 1}; +} + +void applyInterleavedMrope(std::array, 3>& freqs, std::array const& mropeSection) +{ + std::vector merged = freqs[0]; + for (int32_t dim = 1; dim <= 2; ++dim) + { + int32_t const length = mropeSection[static_cast(dim)] * 3; + for (int32_t j = dim; j < length; j += 3) + { + merged[static_cast(j)] = freqs[static_cast(dim)][static_cast(j)]; + } + } + freqs[0] = std::move(merged); +} + +void computeCosmosRotaryEmbeddings(std::vector const& positionIds, int32_t sequenceLength, + CosmosEngineConfig const& config, std::vector& cosOut, std::vector& sinOut) +{ + int32_t const headDim = config.headDim; + int32_t const halfDim = headDim / 2; + cosOut.assign(static_cast(sequenceLength) * static_cast(headDim), 0.F); + sinOut.assign(static_cast(sequenceLength) * static_cast(headDim), 0.F); + + std::vector invFreq(static_cast(halfDim)); + for (int32_t i = 0; i < halfDim; ++i) + { + invFreq[static_cast(i)] + = 1.F / std::pow(config.ropeTheta, static_cast(i * 2) / static_cast(headDim)); + } + + for (int32_t token = 0; token < sequenceLength; ++token) + { + std::array, 3> freqs; + for (int32_t axis = 0; axis < 3; ++axis) + { + freqs[static_cast(axis)].resize(static_cast(halfDim)); + int64_t const pos = positionIds[static_cast(axis) * static_cast(sequenceLength) + + static_cast(token)]; + float const posF = static_cast(pos); + for (int32_t j = 0; j < halfDim; ++j) + { + freqs[static_cast(axis)][static_cast(j)] + = invFreq[static_cast(j)] * posF; + } + } + + applyInterleavedMrope(freqs, config.mropeSection); + auto const& merged = freqs[0]; + + for (int32_t j = 0; j < halfDim; ++j) + { + float const angle = merged[static_cast(j)]; + float const c = std::cos(angle); + float const s = std::sin(angle); + cosOut[static_cast(token) * static_cast(headDim) + static_cast(j)] + = c; + cosOut[static_cast(token) * static_cast(headDim) + static_cast(j) + + static_cast(halfDim)] = c; + sinOut[static_cast(token) * static_cast(headDim) + static_cast(j)] + = s; + sinOut[static_cast(token) * static_cast(headDim) + static_cast(j) + + static_cast(halfDim)] = s; + } + } +} + +void uploadHostFloatToHalfGpu(rt::Tensor& dst, std::vector const& host, cudaStream_t stream) +{ + std::vector converted(host.size()); + for (std::size_t i = 0; i < host.size(); ++i) + { + converted[i] = __float2half(host[i]); + } + CUDA_CHECK(cudaMemcpyAsync( + dst.rawPointer(), converted.data(), converted.size() * sizeof(half), cudaMemcpyHostToDevice, stream)); +} + +std::string formatCosmosUserPrompt(std::string const& prompt, CosmosEngineConfig const& config) +{ + std::string text = prompt; + while (!text.empty() && text.back() == '.') + { + text.pop_back(); + } + + if (config.numFrames == 1) + { + return text + ". This image is of " + std::to_string(config.height) + "x" + std::to_string(config.width) + + " resolution."; + } + + float const duration = static_cast(config.numFrames) / config.fps; + std::ostringstream oss; + oss << text << ". The video is " << std::fixed << std::setprecision(1) << duration << " seconds long and is of " + << static_cast(config.fps) << " FPS." + << ". This video is of " << config.height << "x" << config.width << " resolution."; + return oss.str(); +} + +} // namespace + +std::vector tokenizeCosmosPrompt( + tokenizer::Tokenizer const& tokenizer, std::string const& prompt, CosmosEngineConfig const& config) +{ + std::vector inputIds; + std::string const userPrompt = formatCosmosUserPrompt(prompt, config); + if (config.useChatTemplate) + { + LLMGenerationRequest::Request request; + if (config.useSystemPrompt) + { + Message::MessageContent systemContent; + systemContent.type = "text"; + systemContent.content = config.systemPromptVideo; + + Message systemMessage; + systemMessage.role = "system"; + systemMessage.contents.push_back(std::move(systemContent)); + request.messages.push_back(std::move(systemMessage)); + } + + Message::MessageContent userContent; + userContent.type = "text"; + userContent.content = userPrompt; + + Message userMessage; + userMessage.role = "user"; + userMessage.contents.push_back(std::move(userContent)); + request.messages.push_back(std::move(userMessage)); + + LLMGenerationRequest::FormattedRequest formatted; + if (!tokenizer.applyChatTemplate(request, formatted, true, true, false)) + { + throw std::runtime_error("Cosmos chat template application failed."); + } + + inputIds = tokenizer.encode(formatted.formattedCompleteRequest, false, false); + } + else + { + inputIds = tokenizer.encode(userPrompt, false, false); + } + + if (inputIds.empty()) + { + throw std::runtime_error("Cosmos prompt tokenization produced an empty input_ids sequence."); + } + + inputIds.push_back(tokenizer.getEosId()); + + int32_t const visionStartId = tokenizer.getTokenId(config.visionStartToken); + if (visionStartId < 0) + { + throw std::runtime_error("Cosmos tokenizer is missing vision start token: " + config.visionStartToken); + } + inputIds.push_back(visionStartId); + return inputIds; +} + +std::vector buildCosmosPositionIds( + int32_t undLen, CosmosPackedStatic const& packed, CosmosEngineConfig const& config) +{ + if (packed.visionTokenShape.size() != 3U) + { + throw std::runtime_error("Cosmos packing_static is missing vision_token_shapes"); + } + + auto [textIds, nextOffset] = get3dMropeIdsTextTokens(undLen, 0, config.enableFpsModulation); + int64_t const visionTemporalOffset = nextOffset + static_cast(config.unified3dMropeTemporalModalityMargin); + + auto [visionIds, unusedNext] = get3dMropeIdsVaeTokens(packed.visionTokenShape[0], packed.visionTokenShape[1], + packed.visionTokenShape[2], visionTemporalOffset, config.unified3dMropeResetSpatialIds, config.fps, + config.baseFps, config.temporalCompressionFactor, config.enableFpsModulation); + (void) unusedNext; + + int32_t const genLen = packed.visionTokenShape[0] * packed.visionTokenShape[1] * packed.visionTokenShape[2]; + int32_t const sequenceLength = undLen + genLen; + std::vector positionIds(static_cast(sequenceLength) * 3U); + + for (int32_t axis = 0; axis < 3; ++axis) + { + for (int32_t i = 0; i < undLen; ++i) + { + positionIds[static_cast(axis) * static_cast(sequenceLength) + + static_cast(i)] + = textIds[static_cast(axis) * static_cast(undLen) + + static_cast(i)]; + } + for (int32_t i = 0; i < genLen; ++i) + { + positionIds[static_cast(axis) * static_cast(sequenceLength) + + static_cast(undLen) + static_cast(i)] + = visionIds[static_cast(axis) * static_cast(genLen) + + static_cast(i)]; + } + } + + return positionIds; +} + +bool prepareCosmosTextPhase0(tokenizer::Tokenizer const& tokenizer, EmbeddingData const& embedding, + CosmosEngineConfig const& config, CosmosPackedStatic const& packed, std::string const& prompt, + CosmosTextPhase0& phase0, cudaStream_t stream) +{ + try + { + phase0.inputIds = tokenizeCosmosPrompt(tokenizer, prompt, config); + phase0.undLen = static_cast(phase0.inputIds.size()); + phase0.genLen = packed.visionTokenShape[0] * packed.visionTokenShape[1] * packed.visionTokenShape[2]; + phase0.sequenceLength = phase0.undLen + phase0.genLen; + + if (phase0.undLen != packed.undLen) + { + LOG_WARNING( + "Runtime und_len=%d differs from packing_static und_len=%d; vision indexes must be rebuilt before " + "denoise.", + phase0.undLen, packed.undLen); + } + + auto const positionIds = buildCosmosPositionIds(phase0.undLen, packed, config); + + rt::Tensor hostInputIds({1, phase0.undLen}, rt::DeviceType::kCPU, nvinfer1::DataType::kINT32); + std::memcpy(hostInputIds.rawPointer(), phase0.inputIds.data(), + static_cast(phase0.undLen) * sizeof(int32_t)); + + rt::Tensor gpuInputIds({1, phase0.undLen}, rt::DeviceType::kGPU, nvinfer1::DataType::kINT32); + CUDA_CHECK(cudaMemcpyAsync(gpuInputIds.rawPointer(), hostInputIds.rawPointer(), + static_cast(phase0.undLen) * sizeof(int32_t), cudaMemcpyHostToDevice, stream)); + + rt::Tensor embedOutput({1, phase0.undLen, config.hiddenSize}, rt::DeviceType::kGPU, nvinfer1::DataType::kHALF); + kernel::embeddingLookup(gpuInputIds, embedding.table, embedding.scalesAsOptional(), embedOutput, stream); + + if (!phase0.undSeq.reshape({phase0.undLen, config.hiddenSize})) + { + LOG_ERROR("Failed to reshape und_seq to [%d, %d].", phase0.undLen, config.hiddenSize); + return false; + } + CUDA_CHECK(cudaMemcpyAsync(phase0.undSeq.rawPointer(), embedOutput.rawPointer(), + static_cast(phase0.undLen) * static_cast(config.hiddenSize) * sizeof(half), + cudaMemcpyDeviceToDevice, stream)); + + std::vector cosFull; + std::vector sinFull; + computeCosmosRotaryEmbeddings(positionIds, phase0.sequenceLength, config, cosFull, sinFull); + + if (!phase0.cosUnd.reshape({phase0.undLen, config.headDim}) + || !phase0.sinUnd.reshape({phase0.undLen, config.headDim}) + || !phase0.cosGen.reshape({phase0.genLen, config.headDim}) + || !phase0.sinGen.reshape({phase0.genLen, config.headDim})) + { + LOG_ERROR("Failed to reshape Cosmos rotary tensors."); + return false; + } + + std::vector cosUndHost( + static_cast(phase0.undLen) * static_cast(config.headDim)); + std::vector sinUndHost( + static_cast(phase0.undLen) * static_cast(config.headDim)); + std::vector cosGenHost( + static_cast(phase0.genLen) * static_cast(config.headDim)); + std::vector sinGenHost( + static_cast(phase0.genLen) * static_cast(config.headDim)); + + for (int32_t i = 0; i < phase0.undLen; ++i) + { + std::memcpy(cosUndHost.data() + static_cast(i) * static_cast(config.headDim), + cosFull.data() + static_cast(i) * static_cast(config.headDim), + static_cast(config.headDim) * sizeof(float)); + std::memcpy(sinUndHost.data() + static_cast(i) * static_cast(config.headDim), + sinFull.data() + static_cast(i) * static_cast(config.headDim), + static_cast(config.headDim) * sizeof(float)); + } + for (int32_t i = 0; i < phase0.genLen; ++i) + { + std::size_t const src + = static_cast(phase0.undLen + i) * static_cast(config.headDim); + std::size_t const dst = static_cast(i) * static_cast(config.headDim); + std::memcpy(cosGenHost.data() + dst, cosFull.data() + src, + static_cast(config.headDim) * sizeof(float)); + std::memcpy(sinGenHost.data() + dst, sinFull.data() + src, + static_cast(config.headDim) * sizeof(float)); + } + + uploadHostFloatToHalfGpu(phase0.cosUnd, cosUndHost, stream); + uploadHostFloatToHalfGpu(phase0.sinUnd, sinUndHost, stream); + uploadHostFloatToHalfGpu(phase0.cosGen, cosGenHost, stream); + uploadHostFloatToHalfGpu(phase0.sinGen, sinGenHost, stream); + + LOG_INFO("Cosmos Phase 0 ready: und_len=%d, gen_len=%d, seq_len=%d, und_seq=[%d,%d], rotary head_dim=%d", + phase0.undLen, phase0.genLen, phase0.sequenceLength, phase0.undLen, config.hiddenSize, config.headDim); + return true; + } + catch (std::exception const& e) + { + LOG_ERROR("prepareCosmosTextPhase0 failed: %s", e.what()); + return false; + } +} + +bool seedNoisyVisionLatents(rt::Tensor const& cleanLatents, rt::Tensor& noisyLatents, + std::vector const& noisyFrameIndexes, int32_t seed, float noiseScale, cudaStream_t stream) +{ + if (cleanLatents.getDeviceType() != rt::DeviceType::kGPU) + { + LOG_ERROR("seedNoisyVisionLatents: clean latents must be on GPU."); + return false; + } + if (noisyLatents.getDeviceType() != rt::DeviceType::kGPU) + { + LOG_ERROR("seedNoisyVisionLatents: noisy latents must be on GPU."); + return false; + } + + auto const shape = cleanLatents.getShape(); + if (shape.getNumDims() != 5) + { + LOG_ERROR("seedNoisyVisionLatents: expected rank-5 latents [B,C,T,H,W], got %d dims.", shape.getNumDims()); + return false; + } + if (noisyLatents.getShape() != shape) + { + if (!noisyLatents.reshape(shape)) + { + LOG_ERROR("seedNoisyVisionLatents: failed to reshape noisy latents to %s.", shape.formatString().c_str()); + return false; + } + } + if (cleanLatents.getDataType() != noisyLatents.getDataType()) + { + LOG_ERROR("seedNoisyVisionLatents: clean/noisy latent dtypes must match."); + return false; + } + + nvinfer1::DataType const dtype = cleanLatents.getDataType(); + if (dtype != nvinfer1::DataType::kHALF && dtype != nvinfer1::DataType::kFLOAT) + { + LOG_ERROR("seedNoisyVisionLatents: unsupported latent dtype."); + return false; + } + + int64_t const batch = shape[0]; + int64_t const channels = shape[1]; + int64_t const latentT = shape[2]; + int64_t const height = shape[3]; + int64_t const width = shape[4]; + int64_t const framePlaneElems = height * width; + int64_t const temporalStride = framePlaneElems; + std::size_t const elemSize = (dtype == nvinfer1::DataType::kHALF) ? sizeof(half) : sizeof(float); + std::size_t const tensorBytes = static_cast(shape.volume()) * elemSize; + + CUDA_CHECK(cudaMemcpyAsync( + noisyLatents.rawPointer(), cleanLatents.rawPointer(), tensorBytes, cudaMemcpyDeviceToDevice, stream)); + + std::vector noiseHost(static_cast(shape.volume())); + std::mt19937 generator(static_cast(seed)); + std::normal_distribution dist(0.F, 1.F); + for (float& value : noiseHost) + { + value = dist(generator) * noiseScale; + } + + std::vector planeBytes(static_cast(framePlaneElems) * elemSize); + for (int32_t frameIdx : noisyFrameIndexes) + { + if (frameIdx < 0 || frameIdx >= latentT) + { + LOG_ERROR("seedNoisyVisionLatents: invalid noisy frame index %d (latent T=%ld).", frameIdx, latentT); + return false; + } + + for (int64_t b = 0; b < batch; ++b) + { + for (int64_t c = 0; c < channels; ++c) + { + int64_t const linearOffset = ((b * channels + c) * latentT + frameIdx) * temporalStride; + for (int64_t elem = 0; elem < framePlaneElems; ++elem) + { + float const value = noiseHost[static_cast(linearOffset + elem)]; + if (dtype == nvinfer1::DataType::kHALF) + { + reinterpret_cast(planeBytes.data())[static_cast(elem)] + = __float2half(value); + } + else + { + reinterpret_cast(planeBytes.data())[static_cast(elem)] = value; + } + } + + std::size_t const dstOffset = static_cast(linearOffset) * elemSize; + CUDA_CHECK(cudaMemcpyAsync(static_cast(noisyLatents.rawPointer()) + dstOffset, + planeBytes.data(), planeBytes.size(), cudaMemcpyHostToDevice, stream)); + } + } + } + + LOG_INFO("seedNoisyVisionLatents: cloned clean latents and noised %zu frame(s) with seed=%d.", + noisyFrameIndexes.size(), seed); + return true; +} + +bool seedNoisySoundLatents(rt::Tensor const& cleanLatents, rt::Tensor& noisyLatents, + std::vector const& noisySlotIndexes, int32_t seed, float noiseScale, cudaStream_t stream) +{ + if (cleanLatents.getDeviceType() != rt::DeviceType::kGPU) + { + LOG_ERROR("seedNoisySoundLatents: clean latents must be on GPU."); + return false; + } + if (noisyLatents.getDeviceType() != rt::DeviceType::kGPU) + { + LOG_ERROR("seedNoisySoundLatents: noisy latents must be on GPU."); + return false; + } + + auto const shape = cleanLatents.getShape(); + if (shape.getNumDims() != 3) + { + LOG_ERROR("seedNoisySoundLatents: expected rank-3 latents [B,C,T], got %d dims.", shape.getNumDims()); + return false; + } + if (noisyLatents.getShape() != shape) + { + if (!noisyLatents.reshape(shape)) + { + LOG_ERROR("seedNoisySoundLatents: failed to reshape noisy latents to %s.", shape.formatString().c_str()); + return false; + } + } + if (cleanLatents.getDataType() != noisyLatents.getDataType()) + { + LOG_ERROR("seedNoisySoundLatents: clean/noisy latent dtypes must match."); + return false; + } + + nvinfer1::DataType const dtype = cleanLatents.getDataType(); + if (dtype != nvinfer1::DataType::kHALF && dtype != nvinfer1::DataType::kFLOAT) + { + LOG_ERROR("seedNoisySoundLatents: unsupported latent dtype."); + return false; + } + + int64_t const batch = shape[0]; + int64_t const channels = shape[1]; + int64_t const latentT = shape[2]; + std::size_t const elemSize = (dtype == nvinfer1::DataType::kHALF) ? sizeof(half) : sizeof(float); + std::size_t const tensorBytes = static_cast(shape.volume()) * elemSize; + + CUDA_CHECK(cudaMemcpyAsync( + noisyLatents.rawPointer(), cleanLatents.rawPointer(), tensorBytes, cudaMemcpyDeviceToDevice, stream)); + + std::vector slots = noisySlotIndexes; + if (slots.empty()) + { + slots.reserve(static_cast(latentT)); + for (int64_t t = 0; t < latentT; ++t) + { + slots.push_back(static_cast(t)); + } + } + + std::vector noiseHost(static_cast(shape.volume())); + std::mt19937 generator(static_cast(seed)); + std::normal_distribution dist(0.F, 1.F); + for (float& value : noiseHost) + { + value = dist(generator) * noiseScale; + } + + std::vector slotBytes(static_cast(channels) * elemSize); + for (int32_t slotIdx : slots) + { + if (slotIdx < 0 || slotIdx >= latentT) + { + LOG_ERROR("seedNoisySoundLatents: invalid noisy slot index %d (latent T=%ld).", slotIdx, latentT); + return false; + } + + for (int64_t b = 0; b < batch; ++b) + { + for (int64_t c = 0; c < channels; ++c) + { + int64_t const linearOffset = (b * channels + c) * latentT + slotIdx; + float const value = noiseHost[static_cast(linearOffset)]; + if (dtype == nvinfer1::DataType::kHALF) + { + reinterpret_cast(slotBytes.data())[static_cast(c)] = __float2half(value); + } + else + { + reinterpret_cast(slotBytes.data())[static_cast(c)] = value; + } + } + + for (int64_t c = 0; c < channels; ++c) + { + std::size_t const dstOffset + = static_cast((b * channels + c) * latentT + slotIdx) * elemSize; + CUDA_CHECK(cudaMemcpyAsync(static_cast(noisyLatents.rawPointer()) + dstOffset, + slotBytes.data() + static_cast(c) * elemSize, elemSize, cudaMemcpyHostToDevice, + stream)); + } + } + } + + LOG_INFO("seedNoisySoundLatents: cloned clean latents and noised %zu slot(s) with seed=%d.", slots.size(), seed); + return true; +} + +} // namespace rt +} // namespace trt_edgellm diff --git a/cpp/runtime/wfmRuntimeUtils.cpp b/cpp/runtime/wfmRuntimeUtils.cpp new file mode 100644 index 00000000..7325c297 --- /dev/null +++ b/cpp/runtime/wfmRuntimeUtils.cpp @@ -0,0 +1,303 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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. + */ + +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/wfmRuntimeUtils.h" + +#include "common/logger.h" + +#include +#include + +namespace trt_edgellm +{ +namespace rt +{ +namespace +{ +std::vector readIntVector(nlohmann::json const& value) +{ + std::vector out; + if (!value.is_array()) + { + return out; + } + out.reserve(value.size()); + for (auto const& item : value) + { + out.push_back(item.get()); + } + return out; +} + +bool videoShapeMatchesConfig(VideoBuffer const& video, CosmosEngineConfig const& config) noexcept +{ + if (!video.buffer) + { + LOG_ERROR("WFM request missing pixel buffer."); + return false; + } + + if (video.batch != 1) + { + LOG_ERROR( + "Cosmos Edge runtime currently supports batch=1 (got batch=%lld).", static_cast(video.batch)); + return false; + } + + if (video.channels != 3) + { + LOG_ERROR("Cosmos Edge expects 3 RGB channels (got channels=%lld).", static_cast(video.channels)); + return false; + } + + if (video.numFrames != config.numFrames || video.height != config.height || video.width != config.width) + { + LOG_ERROR("Video shape [1,3,%lld,%lld,%lld] does not match export config [1,3,%d,%d,%d].", + static_cast(video.numFrames), static_cast(video.height), + static_cast(video.width), config.numFrames, config.height, config.width); + return false; + } + + auto const shape = video.buffer->getShape(); + if (shape.getNumDims() == 5) + { + // [B, C, T, H, W] + bool const ok = shape[0] == video.batch && shape[1] == video.channels && shape[2] == video.numFrames + && shape[3] == video.height && shape[4] == video.width; + if (!ok) + { + LOG_ERROR("Pixel tensor shape does not match VideoBuffer metadata."); + return false; + } + } + else + { + LOG_ERROR("Pixel tensor must be 5D [B, C, T, H, W]; got %d dims.", shape.getNumDims()); + return false; + } + + return true; +} + +bool audioShapeMatchesConfig(AudioBuffer const& audio, CosmosEngineConfig const& config) noexcept +{ + if (!audio.buffer) + { + return true; + } + + if (audio.batch != 1) + { + LOG_ERROR( + "Cosmos sound runtime currently supports batch=1 (got batch=%lld).", static_cast(audio.batch)); + return false; + } + + if (audio.sampleRate != config.sampleRate) + { + LOG_ERROR( + "Audio sample rate %d does not match export config sample_rate=%d.", audio.sampleRate, config.sampleRate); + return false; + } + + auto const shape = audio.buffer->getShape(); + if (shape.getNumDims() == 3) + { + bool const ok = shape[0] == audio.batch && shape[2] == audio.numSamples; + if (!ok) + { + LOG_ERROR("Waveform tensor shape does not match AudioBuffer metadata."); + return false; + } + } + else if (shape.getNumDims() == 2) + { + bool const ok = shape[0] == audio.batch && shape[1] == audio.numSamples; + if (!ok) + { + LOG_ERROR("Waveform tensor shape does not match AudioBuffer metadata."); + return false; + } + } + else + { + LOG_ERROR("Waveform tensor must be 2D [B,T] or 3D [B,1,T]; got %d dims.", shape.getNumDims()); + return false; + } + + return true; +} + +} // namespace + +CosmosEngineConfig loadCosmosEngineConfig(std::filesystem::path const& configPath) +{ + std::ifstream configFile(configPath); + if (!configFile) + { + throw std::runtime_error("Failed to open Cosmos config: " + configPath.string()); + } + + nlohmann::json configJson; + configFile >> configJson; + + CosmosEngineConfig config; + config.numFrames = configJson.value("num_frames", 0); + config.height = configJson.value("height", 0); + config.width = configJson.value("width", 0); + config.fps = configJson.value("fps", 24.F); + config.hiddenSize = configJson.value("hidden_size", 0); + config.numInferenceSteps = configJson.value("num_inference_steps", 35); + config.maxBatchSize = configJson.value("max_batch_size", 1); + config.numTrainTimesteps = configJson.value("num_train_timesteps", 1000); + config.flowShift = configJson.value("flow_shift", 5.F); + config.schedulerSolverOrder = configJson.value("scheduler_solver_order", 2); + config.headDim = configJson.value("head_dim", 128); + config.ropeTheta = configJson.value("rope_theta", 1000000.F); + config.enableFpsModulation = configJson.value("enable_fps_modulation", false); + config.unified3dMropeTemporalModalityMargin = configJson.value("unified_3d_mrope_temporal_modality_margin", 15000); + config.unified3dMropeResetSpatialIds = configJson.value("unified_3d_mrope_reset_spatial_ids", true); + config.baseFps = configJson.value("base_fps", 24.F); + config.temporalCompressionFactor = configJson.value("temporal_compression_factor", 4); + config.visionStartToken = configJson.value("vision_start_token", std::string{"<|vision_start|>"}); + config.useChatTemplate = configJson.value("use_chat_template", true); + config.useSystemPrompt = configJson.value("use_system_prompt", true); + config.systemPromptVideo = configJson.value( + "system_prompt_video", std::string{"You are a helpful assistant who will generate videos from a give prompt."}); + config.enableSound = configJson.value("enable_sound", false); + config.soundDim = configJson.value("sound_dim", 64); + config.sampleRate = configJson.value("sample_rate", 48000); + config.soundLatentFps = configJson.value("sound_latent_fps", 25.F); + + if (configJson.contains("mrope_section") && configJson.at("mrope_section").is_array()) + { + for (std::size_t i = 0; i < 3U && i < configJson.at("mrope_section").size(); ++i) + { + config.mropeSection[i] = configJson.at("mrope_section").at(i).get(); + } + } + + if (config.numFrames <= 0 || config.height <= 0 || config.width <= 0 || config.hiddenSize <= 0) + { + throw std::runtime_error("Invalid Cosmos config in " + configPath.string()); + } + + return config; +} + +CosmosPackedStatic loadCosmosPackedStatic(std::filesystem::path const& packingPath) +{ + std::ifstream packingFile(packingPath); + if (!packingFile) + { + throw std::runtime_error("Failed to open Cosmos packing_static: " + packingPath.string()); + } + + nlohmann::json packingJson; + packingFile >> packingJson; + + CosmosPackedStatic packed; + packed.undLen = packingJson.value("und_len", 0); + packed.sequenceLength = packingJson.value("sequence_length", 0); + packed.numNoisyTokens = packingJson.value("num_noisy_tokens", 0); + packed.inputIds = readIntVector(packingJson["input_ids"]); + packed.textIndexes = readIntVector(packingJson["text_indexes"]); + packed.positionIds = readIntVector(packingJson["position_ids"]); + packed.visionSequenceIndexes = readIntVector(packingJson["vision_sequence_indexes"]); + packed.visionMseLossIndexes = readIntVector(packingJson["vision_mse_loss_indexes"]); + packed.visionNoisyFrameIndexes = readIntVector(packingJson["vision_noisy_frame_indexes"]); + packed.visionTokenShape = readIntVector(packingJson["vision_token_shapes"].at(0)); + + if (packingJson.contains("sound_token_shapes") && packingJson.at("sound_token_shapes").is_array() + && !packingJson.at("sound_token_shapes").empty()) + { + packed.soundTokenShape = readIntVector(packingJson["sound_token_shapes"].at(0)); + } + if (packingJson.contains("sound_sequence_indexes")) + { + packed.soundSequenceIndexes = readIntVector(packingJson["sound_sequence_indexes"]); + } + if (packingJson.contains("sound_mse_loss_indexes")) + { + packed.soundMseLossIndexes = readIntVector(packingJson["sound_mse_loss_indexes"]); + } + if (packingJson.contains("sound_noisy_slot_indexes")) + { + packed.soundNoisySlotIndexes = readIntVector(packingJson["sound_noisy_slot_indexes"]); + } + + if (packed.sequenceLength <= 0 || packed.visionTokenShape.size() != 3U) + { + throw std::runtime_error("Invalid packing_static contents in " + packingPath.string()); + } + + return packed; +} + +bool examineWFMRequest( + WFMGenerationRequest const& request, CosmosEngineConfig const& config, CosmosPackedStatic const& packed) noexcept +{ + if (request.prompt.empty()) + { + LOG_ERROR("WFM request prompt is empty."); + return false; + } + + int32_t const steps = request.numInferenceSteps > 0 ? request.numInferenceSteps : config.numInferenceSteps; + if (steps <= 0) + { + LOG_ERROR("WFM request numInferenceSteps must be > 0."); + return false; + } + + if (!videoShapeMatchesConfig(request.pixels, config)) + { + return false; + } + + if (request.inputWaveform.buffer && !audioShapeMatchesConfig(request.inputWaveform, config)) + { + return false; + } + + if (request.generateSound && !config.enableSound) + { + LOG_ERROR("WFM request generateSound=true but config enable_sound=false."); + return false; + } + + if (packed.undLen <= 0 || packed.sequenceLength <= packed.undLen) + { + LOG_ERROR("packing_static has invalid und_len=%d sequence_length=%d.", packed.undLen, packed.sequenceLength); + return false; + } + + LOG_INFO( + "WFM request ok: prompt_len=%zu, pixels=[1,3,%d,%d,%d], generate_sound=%s, steps=%d, und_len=%d, seq_len=%d", + request.prompt.size(), config.numFrames, config.height, config.width, request.generateSound ? "yes" : "no", + steps, packed.undLen, packed.sequenceLength); + + return true; +} + +} // namespace rt +} // namespace trt_edgellm diff --git a/cpp/runtime/wfmRuntimeUtils.h b/cpp/runtime/wfmRuntimeUtils.h new file mode 100644 index 00000000..6f208714 --- /dev/null +++ b/cpp/runtime/wfmRuntimeUtils.h @@ -0,0 +1,230 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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. + */ + +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "common/tensor.h" +#include "runtime/llmRuntimeUtils.h" + +#include +#include +#include +#include +#include + +namespace trt_edgellm +{ +namespace tokenizer +{ +class Tokenizer; +} // namespace tokenizer + +namespace rt +{ + +//! Video tensor in Cosmos layout: [batch, channels, numFrames, height, width]. +struct VideoBuffer +{ + std::shared_ptr buffer; + int64_t batch{1}; + int64_t channels{3}; + int64_t numFrames{0}; + int64_t height{0}; + int64_t width{0}; +}; + +//! Audio waveform in Cosmos AVAE layout: [batch, 1, numSamples] or [batch, numSamples]. +struct AudioBuffer +{ + std::shared_ptr buffer; + int64_t batch{1}; + int32_t sampleRate{48000}; + int64_t numSamples{0}; +}; + +//! Export-time shapes and index tables (from Python ``build_cosmos_packed_static``). +struct CosmosPackedStatic +{ + int32_t undLen{0}; + int32_t sequenceLength{0}; + int32_t numNoisyTokens{0}; + std::vector inputIds; + std::vector textIndexes; + std::vector positionIds; //!< Flattened [3, sequenceLength] row-major + std::vector visionSequenceIndexes; + std::vector visionMseLossIndexes; + std::vector visionNoisyFrameIndexes; + std::vector visionTokenShape; //!< [latent_t, patch_h, patch_w] + + std::vector soundSequenceIndexes; + std::vector soundMseLossIndexes; + std::vector soundNoisySlotIndexes; + std::vector soundTokenShape; //!< [sound_dim, latent_frames] or [latent_frames] +}; + +//! Loaded from ``config.json`` next to the exported Cosmos Edge engines. +struct CosmosEngineConfig +{ + int32_t numFrames{0}; + int32_t height{0}; + int32_t width{0}; + float fps{24.F}; + int32_t hiddenSize{0}; + int32_t numInferenceSteps{35}; + int32_t maxBatchSize{1}; + + int32_t numTrainTimesteps{1000}; + float flowShift{5.F}; + int32_t schedulerSolverOrder{2}; + + int32_t headDim{128}; + float ropeTheta{1000000.F}; + std::array mropeSection{24, 20, 20}; + bool enableFpsModulation{false}; + int32_t unified3dMropeTemporalModalityMargin{15000}; + bool unified3dMropeResetSpatialIds{true}; + float baseFps{24.F}; + int32_t temporalCompressionFactor{4}; + std::string visionStartToken{"<|vision_start|>"}; + bool useChatTemplate{true}; + bool useSystemPrompt{true}; + std::string systemPromptVideo{"You are a helpful assistant who will generate videos from a give prompt."}; + + bool enableSound{false}; + int32_t soundDim{64}; + int32_t sampleRate{48000}; + float soundLatentFps{25.F}; +}; + +//! CPU/GPU tensors produced once per prompt before the denoise loop. +struct CosmosTextPhase0 +{ + std::vector inputIds; + int32_t undLen{0}; + int32_t genLen{0}; + int32_t sequenceLength{0}; + rt::Tensor undSeq; + rt::Tensor cosUnd; + rt::Tensor sinUnd; + rt::Tensor cosGen; + rt::Tensor sinGen; +}; + +struct WFMGenerationRequest +{ + std::string prompt; + VideoBuffer pixels; + AudioBuffer inputWaveform; //!< Optional conditioning waveform (requires audio_encode engine). + bool generateSound{false}; //!< When true, run sound denoise/decode path if engines are present. + int32_t numInferenceSteps{0}; //!< 0 = use value from config.json + int32_t seed{0}; +}; + +struct WFMGenerationResponse +{ + VideoBuffer outputVideo; + AudioBuffer outputWaveform; +}; + +//! Latent tensors updated in-place by the denoise loop (vision required, sound optional). +struct CosmosDenoiseLatents +{ + rt::Tensor* vision{nullptr}; + rt::Tensor* sound{nullptr}; +}; + +//! Load export metadata written by the Python Cosmos Edge export script. +CosmosEngineConfig loadCosmosEngineConfig(std::filesystem::path const& configPath); + +CosmosPackedStatic loadCosmosPackedStatic(std::filesystem::path const& packingPath); + +//! Validate request against loaded export config (batch=1 for now). +bool examineWFMRequest( + WFMGenerationRequest const& request, CosmosEngineConfig const& config, CosmosPackedStatic const& packed) noexcept; + +//! Tokenize prompt and append Cosmos boundary tokens (eos + <|vision_start|>). +std::vector tokenizeCosmosPrompt( + tokenizer::Tokenizer const& tokenizer, std::string const& prompt, CosmosEngineConfig const& config); + +//! Build joint mRoPE position IDs [3, sequence_length] for text + vision tokens. +std::vector buildCosmosPositionIds( + int32_t undLen, CosmosPackedStatic const& packed, CosmosEngineConfig const& config); + +//! Lookup und_seq and build und/gen rotary tensors for the MoT backbone. +bool prepareCosmosTextPhase0(tokenizer::Tokenizer const& tokenizer, EmbeddingData const& embedding, + CosmosEngineConfig const& config, CosmosPackedStatic const& packed, std::string const& prompt, + CosmosTextPhase0& phase0, cudaStream_t stream); + +//! Clone clean VAE latents and replace ``vision_noisy_frame_indexes`` slices with Gaussian noise. +bool seedNoisyVisionLatents(rt::Tensor const& cleanLatents, rt::Tensor& noisyLatents, + std::vector const& noisyFrameIndexes, int32_t seed, float noiseScale, cudaStream_t stream); + +//! Clone clean AVAE latents and replace ``sound_noisy_slot_indexes`` temporal slices with Gaussian noise. +//! When \p noisySlotIndexes is empty, all temporal slots are noised (full generation path). +bool seedNoisySoundLatents(rt::Tensor const& cleanLatents, rt::Tensor& noisyLatents, + std::vector const& noisySlotIndexes, int32_t seed, float noiseScale, cudaStream_t stream); + +//! CPU UniPC flow scheduler for vision latents (matches Cosmos3 Edge ``UniPCMultistepScheduler``). +class CosmosVisionScheduler +{ +public: + explicit CosmosVisionScheduler(CosmosEngineConfig const& config); + + void setTimesteps(int32_t numInferenceSteps); + int32_t getNumInferenceSteps() const noexcept + { + return mNumInferenceSteps; + } + int32_t getStepIndex() const noexcept + { + return mStepIndex; + } + float getTimestepValue() const; + bool stepLatents(rt::Tensor& latents, rt::Tensor const& modelOutput, cudaStream_t stream); + +private: + void resetState(); + std::vector convertModelOutput( + std::vector const& sample, std::vector const& modelOutput) const; + std::vector multistepUniPBhUpdate( + std::vector const& sample, std::vector const& convertedOutput, int32_t order) const; + std::vector multistepUniCBhUpdate(std::vector const& convertedOutput, + std::vector const& lastSample, std::vector const& sample, int32_t order) const; + static void sigmaToAlphaSigma(float sigma, float& alpha, float& sigmaOut) noexcept; + + CosmosEngineConfig mConfig{}; + int32_t mNumInferenceSteps{0}; + int32_t mStepIndex{0}; + int32_t mLowerOrderNums{0}; + std::vector mTimesteps; + std::vector mSigmas; + std::vector> mModelOutputs; + std::vector mTimestepList; + std::vector mLastSample; + int32_t mThisOrder{1}; +}; + +//! Alias: Omni uses independent scheduler copies per modality with the same UniPC config. +using CosmosFlowScheduler = CosmosVisionScheduler; + +} // namespace rt +} // namespace trt_edgellm diff --git a/cpp/runtime/wfmSchedulerUtils.cpp b/cpp/runtime/wfmSchedulerUtils.cpp new file mode 100644 index 00000000..55381bc5 --- /dev/null +++ b/cpp/runtime/wfmSchedulerUtils.cpp @@ -0,0 +1,328 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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. + */ + +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/wfmRuntimeUtils.h" + +#include "common/cudaUtils.h" +#include "common/logger.h" + +#include +#include +#include +#include + +namespace trt_edgellm +{ +namespace rt +{ +namespace +{ + +std::vector tensorToHostFloat(rt::Tensor const& tensor, cudaStream_t stream) +{ + auto const elements = static_cast(tensor.getShape().volume()); + std::vector host(elements); + if (tensor.getDataType() == nvinfer1::DataType::kFLOAT) + { + CUDA_CHECK(cudaMemcpyAsync( + host.data(), tensor.rawPointer(), elements * sizeof(float), cudaMemcpyDeviceToHost, stream)); + } + else + { + std::vector hostHalf(elements); + CUDA_CHECK(cudaMemcpyAsync( + hostHalf.data(), tensor.rawPointer(), elements * sizeof(half), cudaMemcpyDeviceToHost, stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + for (std::size_t i = 0; i < elements; ++i) + { + host[i] = __half2float(hostHalf[static_cast(i)]); + } + return host; + } + CUDA_CHECK(cudaStreamSynchronize(stream)); + return host; +} + +void hostFloatToTensor(std::vector const& host, rt::Tensor& tensor, cudaStream_t stream) +{ + auto const elements = static_cast(tensor.getShape().volume()); + if (host.size() != elements) + { + throw std::runtime_error("CosmosVisionScheduler: host/tensor element count mismatch."); + } + + if (tensor.getDataType() == nvinfer1::DataType::kFLOAT) + { + CUDA_CHECK(cudaMemcpyAsync( + tensor.rawPointer(), host.data(), elements * sizeof(float), cudaMemcpyHostToDevice, stream)); + } + else + { + std::vector hostHalf(elements); + for (std::size_t i = 0; i < elements; ++i) + { + hostHalf[i] = __float2half(host[i]); + } + CUDA_CHECK(cudaMemcpyAsync( + tensor.rawPointer(), hostHalf.data(), elements * sizeof(half), cudaMemcpyHostToDevice, stream)); + } +} + +float expm1fStable(float x) +{ + return std::expm1(x); +} + +} // namespace + +CosmosVisionScheduler::CosmosVisionScheduler(CosmosEngineConfig const& config) + : mConfig(config) +{ +} + +void CosmosVisionScheduler::sigmaToAlphaSigma(float sigma, float& alpha, float& sigmaOut) noexcept +{ + alpha = 1.F - sigma; + sigmaOut = sigma; +} + +void CosmosVisionScheduler::resetState() +{ + mStepIndex = 0; + mLowerOrderNums = 0; + mThisOrder = 1; + mModelOutputs.assign(static_cast(mConfig.schedulerSolverOrder), {}); + mTimestepList.assign(static_cast(mConfig.schedulerSolverOrder), 0.F); + mLastSample.clear(); +} + +void CosmosVisionScheduler::setTimesteps(int32_t numInferenceSteps) +{ + if (numInferenceSteps <= 0) + { + throw std::runtime_error("CosmosVisionScheduler: numInferenceSteps must be positive."); + } + + mNumInferenceSteps = numInferenceSteps; + mTimesteps.resize(static_cast(numInferenceSteps)); + mSigmas.resize(static_cast(numInferenceSteps) + 1U); + + std::vector raw(static_cast(numInferenceSteps) + 1U); + float const end = 1.F / static_cast(mConfig.numTrainTimesteps); + for (int32_t i = 0; i <= numInferenceSteps; ++i) + { + float const frac = static_cast(i) / static_cast(numInferenceSteps); + raw[static_cast(i)] = 1.F + frac * (end - 1.F); + } + + for (int32_t i = 0; i < numInferenceSteps; ++i) + { + float sigma = raw[static_cast(i)]; + sigma = mConfig.flowShift * sigma / (1.F + (mConfig.flowShift - 1.F) * sigma); + if (i == 0 && std::fabs(sigma - 1.F) < 1e-6F) + { + sigma -= 1e-6F; + } + mSigmas[static_cast(i)] = sigma; + mTimesteps[static_cast(i)] = sigma * static_cast(mConfig.numTrainTimesteps); + } + mSigmas[static_cast(numInferenceSteps)] = 0.F; + resetState(); +} + +float CosmosVisionScheduler::getTimestepValue() const +{ + if (mStepIndex < 0 || mStepIndex >= mNumInferenceSteps) + { + throw std::runtime_error("CosmosVisionScheduler: timestep requested outside inference range."); + } + return mTimesteps[static_cast(mStepIndex)]; +} + +std::vector CosmosVisionScheduler::convertModelOutput( + std::vector const& sample, std::vector const& modelOutput) const +{ + float sigmaT = 0.F; + float alphaUnused = 0.F; + sigmaToAlphaSigma(mSigmas[static_cast(mStepIndex)], alphaUnused, sigmaT); + std::vector converted(sample.size()); + for (std::size_t i = 0; i < sample.size(); ++i) + { + converted[i] = sample[i] - sigmaT * modelOutput[i]; + } + return converted; +} + +std::vector CosmosVisionScheduler::multistepUniPBhUpdate( + std::vector const& sample, std::vector const& convertedOutput, int32_t order) const +{ + float sigmaS0 = mSigmas[static_cast(mStepIndex)]; + float sigmaT = mSigmas[static_cast(mStepIndex + 1)]; + float alphaS0 = 0.F; + float alphaT = 0.F; + float sigmaS0T = 0.F; + float sigmaTVal = 0.F; + sigmaToAlphaSigma(sigmaS0, alphaS0, sigmaS0T); + sigmaToAlphaSigma(sigmaT, alphaT, sigmaTVal); + + float const lambdaT = std::log(alphaT) - std::log(sigmaTVal); + float const lambdaS0 = std::log(alphaS0) - std::log(sigmaS0T); + float const h = lambdaT - lambdaS0; + float const hh = -h; + float const hPhi1 = expm1fStable(hh); + float const B_h = expm1fStable(hh); + + std::vector d1s; + std::vector rhosP; + if (order > 1 && mStepIndex > 0 && !mModelOutputs[static_cast(mModelOutputs.size() - 2)].empty()) + { + auto const& m0 = convertedOutput; + auto const& mi = mModelOutputs[static_cast(mModelOutputs.size() - 2)]; + float sigmaSi = mSigmas[static_cast(mStepIndex - 1)]; + float alphaSi = 0.F; + float sigmaSiVal = 0.F; + sigmaToAlphaSigma(sigmaSi, alphaSi, sigmaSiVal); + float const lambdaSi = std::log(alphaSi) - std::log(sigmaSiVal); + float const rk = (lambdaSi - lambdaS0) / h; + d1s.resize(sample.size()); + for (std::size_t i = 0; i < sample.size(); ++i) + { + d1s[i] = (mi[i] - m0[i]) / rk; + } + rhosP = {0.5F}; + } + + std::vector prev(sample.size()); + for (std::size_t i = 0; i < sample.size(); ++i) + { + float predRes = 0.F; + if (!d1s.empty()) + { + predRes = rhosP[0] * d1s[i]; + } + float const xTilde = (sigmaTVal / sigmaS0T) * sample[i] - alphaT * hPhi1 * convertedOutput[i]; + prev[i] = xTilde - alphaT * B_h * predRes; + } + return prev; +} + +std::vector CosmosVisionScheduler::multistepUniCBhUpdate(std::vector const& convertedOutput, + std::vector const& lastSample, std::vector const& sample, int32_t order) const +{ + if (mStepIndex <= 0) + { + return sample; + } + + float sigmaS0 = mSigmas[static_cast(mStepIndex - 1)]; + float sigmaT = mSigmas[static_cast(mStepIndex)]; + float alphaS0 = 0.F; + float alphaT = 0.F; + float sigmaS0T = 0.F; + float sigmaTVal = 0.F; + sigmaToAlphaSigma(sigmaS0, alphaS0, sigmaS0T); + sigmaToAlphaSigma(sigmaT, alphaT, sigmaTVal); + + float const lambdaT = std::log(alphaT) - std::log(sigmaTVal); + float const lambdaS0 = std::log(alphaS0) - std::log(sigmaS0T); + float const h = lambdaT - lambdaS0; + float const hh = -h; + float const hPhi1 = expm1fStable(hh); + float const B_h = expm1fStable(hh); + + auto const& m0 = mModelOutputs.back(); + std::vector corrected(sample.size()); + for (std::size_t i = 0; i < sample.size(); ++i) + { + float const xTilde = (sigmaTVal / sigmaS0T) * lastSample[i] - alphaT * hPhi1 * m0[i]; + corrected[i] = xTilde - alphaT * B_h * (convertedOutput[i] - m0[i]); + } + (void) order; + return corrected; +} + +bool CosmosVisionScheduler::stepLatents(rt::Tensor& latents, rt::Tensor const& modelOutput, cudaStream_t stream) +{ + if (mStepIndex >= mNumInferenceSteps) + { + LOG_ERROR("CosmosVisionScheduler: step called after all inference steps completed."); + return false; + } + if (latents.getShape() != modelOutput.getShape()) + { + LOG_ERROR("CosmosVisionScheduler: latent/model_output shape mismatch."); + return false; + } + + try + { + auto sample = tensorToHostFloat(latents, stream); + auto modelHost = tensorToHostFloat(modelOutput, stream); + auto converted = convertModelOutput(sample, modelHost); + + bool const useCorrector = mStepIndex > 0 && !mLastSample.empty(); + if (useCorrector) + { + sample = multistepUniCBhUpdate(converted, mLastSample, sample, mThisOrder); + } + + for (std::size_t i = 0; i + 1U < mModelOutputs.size(); ++i) + { + mModelOutputs[i] = std::move(mModelOutputs[i + 1]); + mTimestepList[i] = mTimestepList[i + 1]; + } + mModelOutputs.back() = converted; + mTimestepList.back() = getTimestepValue(); + + int32_t const orderCap = mConfig.schedulerSolverOrder; + int32_t thisOrder = orderCap; + if (mConfig.numInferenceSteps > 0) + { + thisOrder = std::min(orderCap, mNumInferenceSteps - mStepIndex); + } + mThisOrder = std::min(thisOrder, mLowerOrderNums + 1); + if (mThisOrder <= 0) + { + mThisOrder = 1; + } + + mLastSample = sample; + sample = multistepUniPBhUpdate(sample, converted, mThisOrder); + + if (mLowerOrderNums < mConfig.schedulerSolverOrder) + { + ++mLowerOrderNums; + } + ++mStepIndex; + + hostFloatToTensor(sample, latents, stream); + return true; + } + catch (std::exception const& e) + { + LOG_ERROR("CosmosVisionScheduler::stepLatents failed: %s", e.what()); + return false; + } +} + +} // namespace rt +} // namespace trt_edgellm diff --git a/cpp/tokenizer/tokenizer.cpp b/cpp/tokenizer/tokenizer.cpp index 9e41341b..79cd9390 100644 --- a/cpp/tokenizer/tokenizer.cpp +++ b/cpp/tokenizer/tokenizer.cpp @@ -1013,6 +1013,8 @@ bool Tokenizer::loadChatTemplate(std::filesystem::path const& chatTemplateFile) mChatTemplate.generationPromptThinking = jsonData.value("generation_prompt_thinking", ""); mChatTemplate.defaultSystemPrompt = jsonData.value("default_system_prompt", mChatTemplate.defaultSystemPrompt); mChatTemplate.trimContent = jsonData.value("trim_content", false); + mChatTemplate.prefixStrategy = jsonData.value("prefix_strategy", ""); + mChatTemplate.maxSeqLen = jsonData.value("max_seq_len", 0); } catch (std::exception const& e) { diff --git a/cpp/tokenizer/tokenizer.h b/cpp/tokenizer/tokenizer.h index 389167ed..c61f6f30 100644 --- a/cpp/tokenizer/tokenizer.h +++ b/cpp/tokenizer/tokenizer.h @@ -68,6 +68,8 @@ struct ChatTemplateConfig std::string generationPromptThinking; //!< Generation prompt with thinking enabled (optional, model-specific) std::string defaultSystemPrompt; //!< Default system prompt bool trimContent{false}; //!< Whether to trim whitespace from message content (matches Jinja | trim) + std::string prefixStrategy; //!< Optional VLA prefix padding strategy (e.g. pi05_compact_prefix) + int32_t maxSeqLen{0}; //!< Optional padded prefix capacity from processed_chat_template.json }; /*! @@ -319,6 +321,18 @@ class Tokenizer return mChatTemplate.defaultSystemPrompt; } + //! Prefix padding strategy from processed_chat_template.json (e.g. pi05_compact_prefix). + std::string const& getPrefixStrategy() const noexcept + { + return mChatTemplate.prefixStrategy; + } + + //! Padded prefix capacity from processed_chat_template.json (0 = use engine max only). + int32_t getMaxSeqLen() const noexcept + { + return mChatTemplate.maxSeqLen; + } + protected: /** * @brief Parse tokenizer.json to extract configuration diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index cb61e7c3..0355a87c 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -11,3 +11,4 @@ add_subdirectory(utils) add_subdirectory(llm) add_subdirectory(multimodal) add_subdirectory(omni) +add_subdirectory(wfm) diff --git a/examples/llm/CMakeLists.txt b/examples/llm/CMakeLists.txt index 0d6766b4..3ba9b749 100644 --- a/examples/llm/CMakeLists.txt +++ b/examples/llm/CMakeLists.txt @@ -24,6 +24,14 @@ target_include_directories( ${CMAKE_SOURCE_DIR}/examples/utils) add_cross_build_link_options(llm_inference) +add_executable(vla_inference vla_inference.cpp) +target_link_libraries(vla_inference PRIVATE edgellmCore exampleUtils + commonLibraryExt) +target_include_directories( + vla_inference PRIVATE ${COMMON_INCLUDE_DIRS} + ${CMAKE_SOURCE_DIR}/examples/utils) +add_cross_build_link_options(vla_inference) + add_executable(llm_stream llm_stream.cpp) target_link_libraries(llm_stream PRIVATE edgellmCore exampleUtils commonLibraryExt) diff --git a/examples/llm/vla_inference.cpp b/examples/llm/vla_inference.cpp new file mode 100644 index 00000000..c05e95ad --- /dev/null +++ b/examples/llm/vla_inference.cpp @@ -0,0 +1,713 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025 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. + */ + +// Dedicated VLA (Vision-Language-Action) inference driver. +// +// Unlike `llm_inference` (which targets the paged-KV EngineExecutor runtime), this driver +// uses `VlaInferenceRuntime`, the LLMEngineRunner + ActionRunner based path that consumes +// the VLA engine binding layout (past_key_values_*, prefix_k/prefix_v) produced by the +// Test export harness for models such as PI0.5. + +#include "common/checkMacros.h" +#include "common/inputLimits.h" +#include "common/trtUtils.h" +#include "profileFormatter.h" // sanitizeUtf8ForJson +#include "profiling/nvtx_wrapper.h" +#include "runtime/audioUtils.h" +#include "runtime/imageUtils.h" +#include "runtime/llmRuntimeUtils.h" +#include "runtime/vlaInferenceRuntime.h" +#include "tokenizer/tokenizer.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace trt_edgellm; +using Json = nlohmann::json; + +enum VlaInferenceOptionId : int +{ + HELP = 900, + INPUT_FILE = 901, + ENGINE_DIR = 902, + MULTIMODAL_ENGINE_DIR = 903, + OUTPUT_FILE = 904, + DEBUG = 905, + DUMP_OUTPUT = 906, + BATCH_SIZE = 907, + MAX_GENERATE_LENGTH = 908, + ACTION_NOISE_SEED = 909 +}; + +struct VlaInferenceArgs +{ + bool help{false}; + std::string engineDir; + std::string multimodalEngineDir{""}; + std::string inputFile; + std::string outputFile{""}; + bool debug{false}; + bool dumpOutput{false}; + int32_t batchSize{-1}; // -1 means use value from input file + int64_t maxGenerateLength{-1}; // -1 means use value from input file + int32_t actionNoiseSeed{-1}; // -1 means leave runtime default +}; + +void printUsage(char const* programName) +{ + std::cerr << "Usage: " << programName + << " [--help] --engineDir= [--multimodalEngineDir=] --inputFile= " + "--outputFile= [--debug] [--dumpOutput] [--batchSize=] " + "[--maxGenerateLength=] [--actionNoiseSeed=]" + << std::endl; + std::cerr << "Options:" << std::endl; + std::cerr << " --help Display this help message" << std::endl; + std::cerr << " --inputFile Path to input JSON file with requests" << std::endl; + std::cerr << " --engineDir Path to LLM engine directory" << std::endl; + std::cerr << " --multimodalEngineDir Path to multimodal (vision) engine directory (optional)" << std::endl; + std::cerr << " --outputFile Path to output JSON file" << std::endl; + std::cerr << " --debug Enable debug logging" << std::endl; + std::cerr << " --dumpOutput Dump inference output to console" << std::endl; + std::cerr << " --batchSize Override batch size from input file" << std::endl; + std::cerr << " --maxGenerateLength Override max generate length from input file" << std::endl; + std::cerr << " --actionNoiseSeed Seed for action diffusion noise initialization (optional)" << std::endl; +} + +bool parseVlaInferenceArgs(VlaInferenceArgs& args, int argc, char* argv[]) +{ + static struct option inferenceOptions[] = {{"help", no_argument, 0, VlaInferenceOptionId::HELP}, + {"inputFile", required_argument, 0, VlaInferenceOptionId::INPUT_FILE}, + {"engineDir", required_argument, 0, VlaInferenceOptionId::ENGINE_DIR}, + {"multimodalEngineDir", required_argument, 0, VlaInferenceOptionId::MULTIMODAL_ENGINE_DIR}, + {"outputFile", required_argument, 0, VlaInferenceOptionId::OUTPUT_FILE}, + {"debug", no_argument, 0, VlaInferenceOptionId::DEBUG}, + {"dumpOutput", no_argument, 0, VlaInferenceOptionId::DUMP_OUTPUT}, + {"batchSize", required_argument, 0, VlaInferenceOptionId::BATCH_SIZE}, + {"maxGenerateLength", required_argument, 0, VlaInferenceOptionId::MAX_GENERATE_LENGTH}, + {"actionNoiseSeed", required_argument, 0, VlaInferenceOptionId::ACTION_NOISE_SEED}, {0, 0, 0, 0}}; + + int opt; + while ((opt = getopt_long(argc, argv, "", inferenceOptions, nullptr)) != -1) + { + switch (opt) + { + case VlaInferenceOptionId::HELP: args.help = true; return true; + case VlaInferenceOptionId::INPUT_FILE: args.inputFile = optarg; break; + case VlaInferenceOptionId::ENGINE_DIR: args.engineDir = optarg; break; + case VlaInferenceOptionId::MULTIMODAL_ENGINE_DIR: args.multimodalEngineDir = optarg; break; + case VlaInferenceOptionId::OUTPUT_FILE: args.outputFile = optarg; break; + case VlaInferenceOptionId::DEBUG: args.debug = true; break; + case VlaInferenceOptionId::DUMP_OUTPUT: args.dumpOutput = true; break; + case VlaInferenceOptionId::BATCH_SIZE: + try + { + args.batchSize = std::stoi(optarg); + if (args.batchSize <= 0) + { + LOG_ERROR("Invalid batchSize value: %s (must be positive)", optarg); + return false; + } + } + catch (std::exception const& e) + { + LOG_ERROR("Invalid batchSize value: %s", optarg); + return false; + } + break; + case VlaInferenceOptionId::MAX_GENERATE_LENGTH: + try + { + args.maxGenerateLength = std::stoll(optarg); + if (args.maxGenerateLength < 0) + { + LOG_ERROR("Invalid maxGenerateLength value: %s (must be non-negative)", optarg); + return false; + } + } + catch (std::exception const& e) + { + LOG_ERROR("Invalid maxGenerateLength value: %s", optarg); + return false; + } + break; + case VlaInferenceOptionId::ACTION_NOISE_SEED: + try + { + args.actionNoiseSeed = std::stoi(optarg); + } + catch (std::exception const& e) + { + LOG_ERROR("Invalid actionNoiseSeed value: %s", optarg); + return false; + } + break; + default: return false; + } + } + + if (args.inputFile.empty()) + { + LOG_ERROR("ERROR: --inputFile is required"); + return false; + } + if (args.engineDir.empty()) + { + LOG_ERROR("ERROR: --engineDir is required"); + return false; + } + if (args.outputFile.empty()) + { + LOG_ERROR("ERROR: --outputFile is required"); + return false; + } + LOG_INFO("args.inputFile: %s", args.inputFile.c_str()); + LOG_INFO("args.engineDir: %s", args.engineDir.c_str()); + if (!args.multimodalEngineDir.empty()) + { + LOG_INFO("args.multimodalEngineDir: %s", args.multimodalEngineDir.c_str()); + } + LOG_INFO("args.outputFile: %s", args.outputFile.c_str()); + + if (args.debug) + { + gLogger.setLevel(nvinfer1::ILogger::Severity::kVERBOSE); + } + else + { + gLogger.setLevel(nvinfer1::ILogger::Severity::kINFO); + } + + return true; +} + +namespace +{ + +std::vector loadRobotStateBin(std::filesystem::path const& path) +{ + std::ifstream file(path, std::ios::binary | std::ios::ate); + check::check(file.is_open(), "Failed to open robot state file: " + path.string()); + auto const fileSize = static_cast(file.tellg()); + check::check( + fileSize % sizeof(float) == 0, "Robot state file size must be a multiple of 4 bytes: " + path.string()); + file.seekg(0); + std::vector values(fileSize / sizeof(float)); + file.read(reinterpret_cast(values.data()), static_cast(fileSize)); + check::check(file.good(), "Failed to read robot state file: " + path.string()); + return values; +} + +std::vector flattenRobotStateJson(nlohmann::json const& stateJson) +{ + std::vector values; + if (stateJson.is_array()) + { + if (!stateJson.empty() && stateJson.front().is_array()) + { + for (auto const& row : stateJson) + { + check::check(row.is_array(), "robot_state nested arrays must contain numeric values"); + for (auto const& value : row) + { + values.push_back(value.get()); + } + } + } + else + { + for (auto const& value : stateJson) + { + values.push_back(value.get()); + } + } + } + else + { + throw std::runtime_error("robot_state must be a JSON array"); + } + return values; +} + +std::pair, std::vector> parseInputFile( + std::filesystem::path const& inputFilePath, int32_t batchSizeOverride = -1, int64_t maxGenerateLengthOverride = -1) +{ + std::vector batchedRequests; + + Json inputData; + std::ifstream inputFileStream(inputFilePath); + check::check(inputFileStream.is_open(), "Failed to open input file: " + inputFilePath.string()); + try + { + inputData = Json::parse(inputFileStream); + inputFileStream.close(); + } + catch (Json::parse_error const& e) + { + throw std::runtime_error( + format::fmtstr("Failed to parse input file %s with error: %s", inputFilePath.string().c_str(), e.what())); + } + + int batchSize = (batchSizeOverride != -1) ? batchSizeOverride : inputData.value("batch_size", 1); + check::check(batchSize > 0, format::fmtstr("Invalid batch_size value: %d (must be positive)", batchSize)); + check::check(batchSize <= limits::security::kReasonableMaxBatchSize, + format::fmtstr("Input rejected: batch_size %d exceeds limit %d. Limit defined in %s.", batchSize, + limits::security::kReasonableMaxBatchSize, limits::kInputLimitsLocation)); + + float temperature = inputData.value("temperature", 1.0f); + float topP = inputData.value("top_p", 0.8f); + int64_t topK = inputData.value("top_k", 50); + int64_t maxGenerateLength + = (maxGenerateLengthOverride != -1) ? maxGenerateLengthOverride : inputData.value("max_generate_length", 256); + check::check(maxGenerateLength >= 0, + format::fmtstr("Invalid max_generate_length value: %lld (must be non-negative)", + static_cast(maxGenerateLength))); + + bool applyChatTemplate = inputData.value("apply_chat_template", true); + bool addGenerationPrompt = inputData.value("add_generation_prompt", true); + bool enableThinking = inputData.value("enable_thinking", false); + + std::unordered_map loraWeightsMap; + if (inputData.contains("available_lora_weights") && inputData["available_lora_weights"].is_object()) + { + auto const& availableLoraWeights = inputData["available_lora_weights"]; + for (auto const& [loraName, loraPath] : availableLoraWeights.items()) + { + check::check(loraPath.is_string(), "LoRA weight path for '" + loraName + "' must be a string"); + check::check(loraWeightsMap.find(loraName) == loraWeightsMap.end(), + "Lora weights with name " + loraName + " already exists"); + loraWeightsMap[loraName] = loraPath.get(); + LOG_INFO("Registered LoRA weights '%s' -> '%s'", loraName.c_str(), loraWeightsMap[loraName].c_str()); + } + } + + if (!(inputData.contains("requests") && inputData["requests"].is_array())) + { + throw std::runtime_error("'requests' array not found in input file"); + } + + auto& requestsArray = inputData["requests"]; + size_t numRequests = requestsArray.size(); + + for (size_t startIdx = 0; startIdx < numRequests; startIdx += batchSize) + { + rt::LLMGenerationRequest batchRequest; + batchRequest.temperature = temperature; + batchRequest.topP = topP; + batchRequest.topK = topK; + batchRequest.maxGenerateLength = maxGenerateLength; + batchRequest.applyChatTemplate = applyChatTemplate; + batchRequest.addGenerationPrompt = addGenerationPrompt; + batchRequest.enableThinking = enableThinking; + if (inputData.contains("embodiment_id")) + { + batchRequest.embodimentId = inputData["embodiment_id"].get(); + } + if (inputData.contains("action_batch_size")) + { + batchRequest.actionBatchSize = inputData["action_batch_size"].get(); + } + + std::string batchLoraWeightsName = ""; + bool firstInBatch = true; + + size_t endIdx = std::min(startIdx + batchSize, numRequests); + for (size_t requestIdx = startIdx; requestIdx < endIdx; ++requestIdx) + { + auto const& requestItem = requestsArray[requestIdx]; + check::check(requestItem.is_object(), "Each request must be an object with 'messages' key"); + + bool saveSystemPromptKVCache = requestItem.value("save_system_prompt_kv_cache", false); + if (saveSystemPromptKVCache) + { + batchRequest.saveSystemPromptKVCache = true; + } + + check::check(requestItem.contains("messages") && requestItem["messages"].is_array(), + "Each request object must contain a 'messages' array"); + + auto const& messagesArray = requestItem["messages"]; + + std::string requestLoraName = ""; + if (requestItem.contains("lora_name") && !requestItem["lora_name"].is_null()) + { + requestLoraName = requestItem["lora_name"].get(); + check::check(requestLoraName.empty() || loraWeightsMap.find(requestLoraName) != loraWeightsMap.end(), + "LoRA name '" + requestLoraName + "' not found in available_lora_weights"); + } + + if (firstInBatch) + { + batchLoraWeightsName = requestLoraName; + firstInBatch = false; + } + else + { + check::check(requestLoraName == batchLoraWeightsName, + "Different LoRA weights within the same batch are not supported"); + } + + std::vector chatMessages; + std::vector imageBuffers; + std::vector audioBuffers; + std::optional> requestPastTrajectory; + std::vector requestRobotState; + std::optional requestEmbodimentId; + + if (requestItem.contains("robot_state")) + { + if (requestItem["robot_state"].is_string()) + { + requestRobotState = loadRobotStateBin(requestItem["robot_state"].get()); + } + else + { + requestRobotState = flattenRobotStateJson(requestItem["robot_state"]); + } + } + if (requestItem.contains("embodiment_id")) + { + requestEmbodimentId = requestItem["embodiment_id"].get(); + } + + check::check(messagesArray.size() <= limits::security::kMaxMessagesPerRequest, + format::fmtstr("Input rejected: too many messages in request %zu: %zu (max: %zu). Limit defined in %s.", + requestIdx, messagesArray.size(), limits::security::kMaxMessagesPerRequest, + limits::kInputLimitsLocation)); + + for (auto const& messageJson : messagesArray) + { + check::check(messageJson.contains("role") && messageJson.contains("content"), + "Each message must have 'role' and 'content' fields"); + + rt::Message chatMsg; + chatMsg.role = messageJson["role"].get(); + + auto const& contentJson = messageJson["content"]; + + if (contentJson.is_string()) + { + std::string const& contentStr = contentJson.get(); + check::check(contentStr.size() <= limits::security::kMaxMessageContentSizeBytes, + format::fmtstr( + "Input rejected: message content too large in request %zu: %zu bytes (max: %zu). " + "Limit defined in %s.", + requestIdx, contentStr.size(), limits::security::kMaxMessageContentSizeBytes, + limits::kInputLimitsLocation)); + + rt::Message::MessageContent msgContent; + msgContent.type = "text"; + msgContent.content = contentStr; + chatMsg.contents.push_back(msgContent); + } + else if (contentJson.is_array()) + { + check::check(contentJson.size() <= limits::security::kMaxContentItemsPerMessage, + format::fmtstr("Input rejected: too many content items in message %zu: %zu (max: %zu). " + "Limit defined in %s.", + requestIdx, contentJson.size(), limits::security::kMaxContentItemsPerMessage, + limits::kInputLimitsLocation)); + + for (auto const& contentItemJson : contentJson) + { + check::check(contentItemJson.contains("type"), "Each content item must have a 'type' field"); + + rt::Message::MessageContent msgContent; + msgContent.type = contentItemJson["type"].get(); + + if (msgContent.type == "text") + { + std::string const& textContent = contentItemJson["text"].get(); + check::check(textContent.size() <= limits::security::kMaxMessageContentSizeBytes, + format::fmtstr( + "Input rejected: message content too large in request %zu: %zu bytes (max: %zu). " + "Limit defined in %s.", + requestIdx, textContent.size(), limits::security::kMaxMessageContentSizeBytes, + limits::kInputLimitsLocation)); + msgContent.content = textContent; + } + else if (msgContent.type == "image") + { + msgContent.content = contentItemJson["image"].get(); + auto image = rt::imageUtils::loadImageFromFile(msgContent.content); + if (image.buffer != nullptr) + { + imageBuffers.push_back(std::move(image)); + } + } + else if (msgContent.type == "trajectory") + { + check::check( + contentItemJson.contains("trajectory") && contentItemJson["trajectory"].is_array(), + "Content type 'trajectory' must have a 'trajectory' array of [x,y,z] points"); + std::vector traj; + for (auto const& pt : contentItemJson["trajectory"]) + { + check::check(pt.is_array() && pt.size() == 3, + "Each trajectory point must be a length-3 array [x, y, z]"); + traj.emplace_back(pt[0].get(), pt[1].get(), pt[2].get()); + } + requestPastTrajectory = std::move(traj); + } + else if (msgContent.type == "state") + { + if (contentItemJson.contains("file")) + { + requestRobotState = loadRobotStateBin(contentItemJson["file"].get()); + } + else if (contentItemJson.contains("values")) + { + requestRobotState = flattenRobotStateJson(contentItemJson["values"]); + } + else + { + throw std::runtime_error("Content type 'state' requires 'file' or 'values'"); + } + } + else if (msgContent.type == "embodiment_id") + { + requestEmbodimentId = contentItemJson["embodiment_id"].get(); + } + else + { + throw std::runtime_error( + format::fmtstr("Content type must be 'text', 'image', 'trajectory', 'state', or " + "'embodiment_id', but got: %s", + msgContent.type.c_str())); + } + + chatMsg.contents.push_back(msgContent); + } + } + else + { + throw std::runtime_error("Message content must be a string or an array"); + } + + chatMessages.push_back(chatMsg); + } + + rt::LLMGenerationRequest::Request request; + request.messages = std::move(chatMessages); + request.imageBuffers = std::move(imageBuffers); + request.audioBuffers = std::move(audioBuffers); + request.pastTrajectory = std::move(requestPastTrajectory); + request.robotState = std::move(requestRobotState); + request.embodimentId = requestEmbodimentId; + batchRequest.requests.push_back(std::move(request)); + } + + if (!batchLoraWeightsName.empty()) + { + batchRequest.loraWeightsName = batchLoraWeightsName; + } + + batchedRequests.push_back(std::move(batchRequest)); + } + + return std::make_pair(std::move(loraWeightsMap), std::move(batchedRequests)); +} + +} // namespace + +int main(int argc, char* argv[]) +{ + NVTX_SCOPED_RANGE(nvtx_main, "vla_inference"); + VlaInferenceArgs args; + if (!parseVlaInferenceArgs(args, argc, argv)) + { + printUsage(argv[0]); + return EXIT_FAILURE; + } + if (args.help) + { + printUsage(argv[0]); + return EXIT_SUCCESS; + } + + auto pluginHandles = loadEdgellmPluginLib(); + + std::unordered_map loraWeightsMap; + std::vector batchedRequests; + try + { + std::tie(loraWeightsMap, batchedRequests) + = parseInputFile(args.inputFile, args.batchSize, args.maxGenerateLength); + LOG_INFO("Successfully parsed %zu LoRA weights from input file.", loraWeightsMap.size()); + LOG_INFO("Successfully parsed %zu batches of requests from input file.", batchedRequests.size()); + } + catch (std::exception const& e) + { + LOG_ERROR("Failed to parse input file: %s", e.what()); + return EXIT_FAILURE; + } + + if (batchedRequests.empty()) + { + LOG_ERROR("No valid requests found in input file."); + return EXIT_FAILURE; + } + + cudaStream_t stream; + CUDA_CHECK(cudaStreamCreate(&stream)); + + std::unique_ptr runtime{nullptr}; + try + { + runtime = std::make_unique( + args.engineDir, args.multimodalEngineDir, loraWeightsMap, stream); + } + catch (std::exception const& e) + { + LOG_ERROR("Failed to initialize VlaInferenceRuntime: %s", e.what()); + return EXIT_FAILURE; + } + + if (args.actionNoiseSeed >= 0) + { + runtime->setActionNoiseSeed(args.actionNoiseSeed); + } + + if (!runtime->captureDecodingCUDAGraph(stream)) + { + LOG_WARNING("Failed to capture CUDA graph for decoding usage, proceeding with normal engine execution."); + } + + nlohmann::json outputData; + outputData["input_file"] = args.inputFile; + outputData["responses"] = nlohmann::json::array(); + + bool hasFailedRequest = false; + std::string const errorMessage = "TensorRT Edge LLM cannot handle this request. Fails."; + size_t failedCount = 0; + + LOG_INFO("Processing %zu batched requests...", batchedRequests.size()); + for (size_t requestIdx = 0; requestIdx < batchedRequests.size(); ++requestIdx) + { + auto& request = batchedRequests[requestIdx]; + rt::LLMGenerationResponse response; + + bool requestStatus = runtime->handleRequest(request, response, stream); + + if (requestStatus) + { + if (args.dumpOutput) + { + for (size_t batchIdx = 0; batchIdx < response.outputTexts.size(); ++batchIdx) + { + LOG_INFO("Response for request %zu batch %zu: %s", requestIdx, batchIdx, + response.outputTexts[batchIdx].c_str()); + } + } + } + else + { + hasFailedRequest = true; + failedCount++; + LOG_ERROR("*** FAILED *** Request %zu failed to process!", requestIdx); + } + + for (size_t batchIdx = 0; batchIdx < request.requests.size(); ++batchIdx) + { + nlohmann::json responseJson; + std::string outputText = (requestStatus && batchIdx < response.outputTexts.size()) + ? response.outputTexts[batchIdx] + : errorMessage; + responseJson["output_text"] = sanitizeUtf8ForJson(outputText); + responseJson["request_idx"] = requestIdx; + responseJson["batch_idx"] = batchIdx; + + nlohmann::json messagesJson = nlohmann::json::array(); + for (auto const& msg : request.requests[batchIdx].messages) + { + nlohmann::json msgJson; + msgJson["role"] = msg.role; + msgJson["content"] = nlohmann::json::array(); + for (auto const& content : msg.contents) + { + nlohmann::json contentJson; + contentJson["type"] = content.type; + if (content.type == "text") + { + contentJson["text"] = content.content; + } + else if (content.type == "image") + { + contentJson["image"] = content.content; + } + msgJson["content"].push_back(contentJson); + } + messagesJson.push_back(msgJson); + } + responseJson["messages"] = messagesJson; + + if (requestStatus && batchIdx < response.outputActions.size() && !response.outputActions[batchIdx].empty()) + { + responseJson["actions"] = response.outputActions[batchIdx]; + } + if (requestStatus && batchIdx < response.outputTrajectories.size() + && !response.outputTrajectories[batchIdx].empty()) + { + nlohmann::json trajJson = nlohmann::json::array(); + for (auto const& pt : response.outputTrajectories[batchIdx]) + { + trajJson.push_back(nlohmann::json::array({pt.first, pt.second})); + } + responseJson["trajectory"] = std::move(trajJson); + } + outputData["responses"].push_back(responseJson); + } + } + + LOG_INFO("Processing complete: %zu/%zu batched requests successful", batchedRequests.size() - failedCount, + batchedRequests.size()); + if (failedCount > 0) + { + LOG_ERROR("*** %zu BATCHED REQUESTS FAILED ***", failedCount); + } + + try + { + std::ofstream outputFile(args.outputFile); + if (outputFile.is_open()) + { + outputFile << outputData.dump(4); + outputFile.close(); + LOG_INFO("All responses exported to: %s", args.outputFile.c_str()); + } + else + { + LOG_ERROR("Failed to open output file: %s", args.outputFile.c_str()); + return EXIT_FAILURE; + } + } + catch (std::exception const& e) + { + LOG_ERROR("Failed to write output file: %s", e.what()); + return EXIT_FAILURE; + } + + return hasFailedRequest ? EXIT_FAILURE : EXIT_SUCCESS; +} diff --git a/examples/wfm/CMakeLists.txt b/examples/wfm/CMakeLists.txt new file mode 100644 index 00000000..623f550a --- /dev/null +++ b/examples/wfm/CMakeLists.txt @@ -0,0 +1,11 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & +# AFFILIATES. All rights reserved. SPDX-License-Identifier: +# LicenseRef-NvidiaProprietary + +add_executable(wfm_inference wfm_inference.cpp) +target_link_libraries(wfm_inference PRIVATE edgellmCore exampleUtils + commonLibraryExt) +target_include_directories( + wfm_inference PRIVATE ${COMMON_INCLUDE_DIRS} + ${CMAKE_SOURCE_DIR}/examples/utils) +add_cross_build_link_options(wfm_inference) diff --git a/examples/wfm/wfm_inference.cpp b/examples/wfm/wfm_inference.cpp new file mode 100644 index 00000000..ebf36d4a --- /dev/null +++ b/examples/wfm/wfm_inference.cpp @@ -0,0 +1,889 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-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 "common/checkMacros.h" +#include "common/inputLimits.h" +#include "common/trtUtils.h" +#include "memoryMonitor.h" +#include "profileFormatter.h" +#include "profiling/metrics.h" +#include "profiling/nvtx_wrapper.h" +#include "profiling/timer.h" +#include "runtime/wfmInferenceRuntime.h" +#include "runtime/wfmRuntimeUtils.h" +#include "tokenizer/tokenizer.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace trt_edgellm; +using Json = nlohmann::json; + +/* + * This executable is the command-line "front end" for WFM inference. + * + * It does not implement the neural network algorithms itself. Instead, it: + * 1. Reads command-line options and a JSON file describing one or more requests. + * 2. Loads input pixels/audio into GPU tensors (or creates random input pixels). + * 3. Calls WFMInferenceRuntime::handleRequest(), where encode/denoise/decode occur. + * 4. Copies requested outputs back from the GPU and writes result metadata as JSON. + * + * Most failures in this file concern invalid input, missing files, CUDA transfers, + * or output writing. Model-specific execution is implemented under cpp/runtime/. + */ + +//! Numeric identifiers returned by getopt_long() for the supported CLI options. +//! +//! Values begin at 900 so that they do not collide with ordinary one-character +//! options such as 'h'. This program uses long options only (for example, --help). +enum WfmInferenceOptionId : int +{ + HELP = 900, + INPUT_FILE = 901, + ENGINE_DIR = 902, + OUTPUT_FILE = 904, + DEBUG = 905, + DUMP_PROFILE = 906, + PROFILE_OUTPUT_FILE = 907, + WARMUP = 908, + DUMP_OUTPUT = 909, + NUM_INFERENCE_STEPS = 910, + SEED = 911 +}; + +//! Values collected from the command line. +//! +//! The brace initializers are defaults. For example, if --warmup is omitted, +//! warmup remains 0; if --seed is omitted, seed remains 42. +struct WfmInferenceArgs +{ + bool help{false}; + std::string engineDir; + std::string inputFile; + std::string outputFile; + std::string profileOutputFile; + bool debug{false}; + bool dumpProfile{false}; + int32_t warmup{0}; + bool dumpOutput{false}; + int32_t numInferenceSteps{-1}; //!< -1 = use value from input file / config.json + int32_t seed{42}; +}; + +//! Default values read from the top level of the input JSON file. +//! +//! An individual request can override these values. Zero means "not specified" +//! for both fields in the resolution helpers below. +struct WfmInputGlobals +{ + int32_t numInferenceSteps{0}; + int32_t seed{0}; +}; + +//! Lightweight CPU-side description of one item in the JSON "requests" array. +//! +//! This struct contains strings and scalar settings only. buildWfmRequest() later +//! converts it into a WFMGenerationRequest containing actual GPU tensors. +struct WfmRequestSpec +{ + std::string prompt; + bool generateSound{false}; + int32_t numInferenceSteps{0}; + int32_t seed{0}; + std::string pixelsFile; + std::string waveformFile; + std::string outputVideoFile; + std::string outputWaveformFile; +}; + +//! Print command-line syntax and a description of every accepted option. +//! +//! @param programName The executable name, normally argv[0]. +//! @note Help is written to stderr so it is visible alongside validation errors. +void printUsage(char const* programName) +{ + std::cerr << "Usage: " << programName + << " [--help] [--engineDir=] [--inputFile=] [--outputFile=] [--dumpProfile] " + "[--profileOutputFile=] [--warmup=] [--debug] [--dumpOutput] " + "[--numInferenceSteps=] [--seed=]" + << std::endl; + std::cerr << "Cosmos WFM video generation inference from an exported engine bundle." << std::endl; + std::cerr << "Options:" << std::endl; + std::cerr << " --help Display this help message" << std::endl; + std::cerr << " --inputFile Path to input JSON file with requests (required)" << std::endl; + std::cerr << " --engineDir Path to exported WFM engine directory (required)" << std::endl; + std::cerr << " --outputFile Path to output JSON file (required)" << std::endl; + std::cerr << " --dumpProfile Dump profiling summary to console" << std::endl; + std::cerr << " --profileOutputFile Path to profile JSON output file (optional)" << std::endl; + std::cerr << " --warmup Number of warmup runs using the first request (default: 0)" << std::endl; + std::cerr << " --debug Enable debug logging" << std::endl; + std::cerr << " --dumpOutput Dump inference output shapes to console" << std::endl; + std::cerr << " --numInferenceSteps Override num_inference_steps from input file" << std::endl; + std::cerr << " --seed Default random seed when not set in input JSON (default: 42)" + << std::endl; +} + +//! Parse and validate command-line options. +//! +//! getopt_long() examines argv and returns one WfmInferenceOptionId at a time. +//! Options containing numbers arrive as text, so std::stoi() converts them. +//! +//! @param[out] args Receives all parsed values. +//! @param argc Number of command-line arguments. +//! @param argv Array of argument strings. +//! @return true if parsing succeeds (including --help), otherwise false. +//! @note This also selects INFO or VERBOSE logging after validation. +bool parseWfmInferenceArgs( + WfmInferenceArgs& args, int argc, char* argv[]) // NOLINT(readability-function-cognitive-complexity) +{ + // Each entry maps a long option such as "--engineDir" to an enum value. + // required_argument means the option must be followed by a value. + static struct option inferenceOptions[] = {{"help", no_argument, 0, WfmInferenceOptionId::HELP}, + {"inputFile", required_argument, 0, WfmInferenceOptionId::INPUT_FILE}, + {"engineDir", required_argument, 0, WfmInferenceOptionId::ENGINE_DIR}, + {"outputFile", required_argument, 0, WfmInferenceOptionId::OUTPUT_FILE}, + {"debug", no_argument, 0, WfmInferenceOptionId::DEBUG}, + {"dumpProfile", no_argument, 0, WfmInferenceOptionId::DUMP_PROFILE}, + {"profileOutputFile", required_argument, 0, WfmInferenceOptionId::PROFILE_OUTPUT_FILE}, + {"warmup", required_argument, 0, WfmInferenceOptionId::WARMUP}, + {"dumpOutput", no_argument, 0, WfmInferenceOptionId::DUMP_OUTPUT}, + {"numInferenceSteps", required_argument, 0, WfmInferenceOptionId::NUM_INFERENCE_STEPS}, + {"seed", required_argument, 0, WfmInferenceOptionId::SEED}, {0, 0, 0, 0}}; + + int opt = 0; + // getopt_long() returns -1 after it has consumed all command-line options. + while ((opt = getopt_long(argc, argv, "", inferenceOptions, nullptr)) != -1) + { + switch (opt) + { + case WfmInferenceOptionId::HELP: args.help = true; return true; + case WfmInferenceOptionId::INPUT_FILE: args.inputFile = optarg; break; + case WfmInferenceOptionId::ENGINE_DIR: args.engineDir = optarg; break; + case WfmInferenceOptionId::OUTPUT_FILE: args.outputFile = optarg; break; + case WfmInferenceOptionId::DEBUG: args.debug = true; break; + case WfmInferenceOptionId::DUMP_PROFILE: args.dumpProfile = true; break; + case WfmInferenceOptionId::PROFILE_OUTPUT_FILE: args.profileOutputFile = optarg; break; + case WfmInferenceOptionId::DUMP_OUTPUT: args.dumpOutput = true; break; + case WfmInferenceOptionId::WARMUP: + try + { + args.warmup = std::stoi(optarg); + if (args.warmup < 0) + { + LOG_ERROR("Invalid warmup value: %s (must be non-negative)", optarg); + return false; + } + } + catch (std::exception const&) + { + LOG_ERROR("Invalid warmup value: %s", optarg); + return false; + } + break; + case WfmInferenceOptionId::NUM_INFERENCE_STEPS: + try + { + args.numInferenceSteps = std::stoi(optarg); + if (args.numInferenceSteps <= 0) + { + LOG_ERROR("Invalid numInferenceSteps value: %s (must be positive)", optarg); + return false; + } + } + catch (std::exception const&) + { + LOG_ERROR("Invalid numInferenceSteps value: %s", optarg); + return false; + } + break; + case WfmInferenceOptionId::SEED: + try + { + args.seed = std::stoi(optarg); + } + catch (std::exception const&) + { + LOG_ERROR("Invalid seed value: %s", optarg); + return false; + } + break; + default: return false; + } + } + + // These three paths are necessary for every non-help invocation. + if (args.inputFile.empty()) + { + LOG_ERROR("ERROR: --inputFile is required"); + return false; + } + if (args.engineDir.empty()) + { + LOG_ERROR("ERROR: --engineDir is required"); + return false; + } + if (args.outputFile.empty()) + { + LOG_ERROR("ERROR: --outputFile is required"); + return false; + } + + if (args.debug) + { + gLogger.setLevel(nvinfer1::ILogger::Severity::kVERBOSE); + } + else + { + gLogger.setLevel(nvinfer1::ILogger::Severity::kINFO); + } + + return true; +} + +namespace +{ + +//! Read an entire raw FP16 binary file into CPU memory. +//! +//! A `half` occupies two bytes. The file has no header or shape metadata, so the +//! caller must know the intended tensor shape and optionally supply its element +//! count. This is not a PNG, MP4, WAV, or other container format. +//! +//! @param path File to read. +//! @param expectedElements Required number of FP16 values, or 0 to accept any size. +//! @return A CPU vector containing the file's FP16 values. +//! @throws std::runtime_error (through check::check) if validation or reading fails. +std::vector loadFp16BinaryFile(std::filesystem::path const& path, std::size_t expectedElements) +{ + // ios::ate initially positions the read cursor at the end, allowing tellg() + // to report the file's byte size without reading the file twice. + std::ifstream file(path, std::ios::binary | std::ios::ate); + check::check(file.is_open(), "Failed to open binary file: " + path.string()); + auto const fileSize = static_cast(file.tellg()); + check::check( + fileSize % sizeof(half) == 0, "Binary file size must be a multiple of 2 bytes (fp16): " + path.string()); + std::size_t const numElements = fileSize / sizeof(half); + check::check(expectedElements == 0 || numElements == expectedElements, + format::fmtstr("Binary file %s has %zu fp16 elements, expected %zu", path.string().c_str(), numElements, + expectedElements)); + file.seekg(0); + std::vector values(numElements); + file.read(reinterpret_cast(values.data()), static_cast(fileSize)); + check::check(file.good(), "Failed to read binary file: " + path.string()); + return values; +} + +//! Write CPU FP16 values as a raw binary file. +//! +//! @param path Destination path. +//! @param data Pointer to the first FP16 value. +//! @param numElements Number of values to write. +//! @throws std::runtime_error (through check::check) if opening or writing fails. +void saveFp16BinaryFile(std::filesystem::path const& path, half const* data, std::size_t numElements) +{ + std::ofstream file(path, std::ios::binary); + check::check(file.is_open(), "Failed to open output binary file: " + path.string()); + file.write(reinterpret_cast(data), static_cast(numElements * sizeof(half))); + check::check(file.good(), "Failed to write output binary file: " + path.string()); +} + +//! Create reproducible random input pixels and copy them to a GPU tensor. +//! +//! The tensor uses NCTHW order: +//! N = batch (1), C = color channels (3), T = frames, H = height, W = width. +//! Values are sampled uniformly from [-1, 1], which is the model's normalized +//! pixel range. Supplying the same seed and configuration produces the same data. +//! +//! @param config Engine dimensions used to determine the tensor shape. +//! @param seed Seed for the CPU pseudo-random number generator. +//! @param stream CUDA stream used for the host-to-device copy. +//! @return Shared ownership of an FP16 tensor allocated on the GPU. +std::shared_ptr makeRandomPixels(rt::CosmosEngineConfig const& config, int32_t seed, cudaStream_t stream) +{ + // Tensor allocates device memory because DeviceType::kGPU is requested. + auto pixels = std::make_shared(rt::Coords({1, 3, config.numFrames, config.height, config.width}), + rt::DeviceType::kGPU, nvinfer1::DataType::kHALF, "wfm_inference::input_pixels"); + + std::vector host(static_cast(pixels->getShape().volume())); + std::mt19937 rng(static_cast(seed)); + std::uniform_real_distribution dist(-1.F, 1.F); + for (auto& value : host) + { + value = __float2half(dist(rng)); + } + + // cudaMemcpyAsync schedules the copy. Synchronization keeps the temporary + // CPU vector alive until the GPU has finished reading from it. + CUDA_CHECK( + cudaMemcpyAsync(pixels->rawPointer(), host.data(), host.size() * sizeof(half), cudaMemcpyHostToDevice, stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + return pixels; +} + +//! Load a raw FP16 pixel tensor from disk and copy it to the GPU. +//! +//! @param path Raw FP16 input file. +//! @param config Supplies the required frame count, height, and width. +//! @param stream CUDA stream used for the host-to-device copy. +//! @return GPU tensor with shape [1, 3, frames, height, width]. +//! @note The binary file has no shape metadata; its element order must already +//! match the model's expected NCTHW layout. +std::shared_ptr loadPixelsFromFile( + std::filesystem::path const& path, rt::CosmosEngineConfig const& config, cudaStream_t stream) +{ + std::size_t const expectedElements = static_cast(config.numFrames) * config.height * config.width * 3U; + auto host = loadFp16BinaryFile(path, expectedElements); + auto pixels = std::make_shared(rt::Coords({1, 3, config.numFrames, config.height, config.width}), + rt::DeviceType::kGPU, nvinfer1::DataType::kHALF, "wfm_inference::input_pixels"); + CUDA_CHECK( + cudaMemcpyAsync(pixels->rawPointer(), host.data(), host.size() * sizeof(half), cudaMemcpyHostToDevice, stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + return pixels; +} + +//! Load a mono FP16 waveform from disk and copy it to the GPU. +//! +//! Unlike the pixel loader, the waveform loader accepts any non-empty length. +//! The number of samples is inferred directly from the file size. +//! +//! @param path Raw FP16 waveform file (not a WAV container). +//! @param config Engine configuration; currently unused by this helper. +//! @param stream CUDA stream used for the host-to-device copy. +//! @param[out] numSamples Receives the inferred waveform length. +//! @return GPU tensor with shape [batch=1, channel=1, samples]. +std::shared_ptr loadWaveformFromFile( + std::filesystem::path const& path, rt::CosmosEngineConfig const& config, cudaStream_t stream, int64_t& numSamples) +{ + auto host = loadFp16BinaryFile(path, 0); + numSamples = static_cast(host.size()); + check::check(numSamples > 0, "Waveform file is empty: " + path.string()); + + auto waveform = std::make_shared(rt::Coords({1, 1, numSamples}), rt::DeviceType::kGPU, + nvinfer1::DataType::kHALF, "wfm_inference::input_waveform"); + CUDA_CHECK(cudaMemcpyAsync( + waveform->rawPointer(), host.data(), host.size() * sizeof(half), cudaMemcpyHostToDevice, stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + return waveform; +} + +//! Copy an FP16 tensor from GPU memory and save it as raw binary. +//! +//! @param path Destination file. +//! @param tensor Source GPU tensor. Other data types are rejected. +//! @param stream CUDA stream used for the device-to-host copy. +void saveTensorToFp16File(std::filesystem::path const& path, rt::Tensor const& tensor, cudaStream_t stream) +{ + check::check(tensor.getDataType() == nvinfer1::DataType::kHALF, "saveTensorToFp16File only supports fp16 tensors"); + std::size_t const numElements = static_cast(tensor.getShape().volume()); + std::vector host(numElements); + CUDA_CHECK( + cudaMemcpyAsync(host.data(), tensor.rawPointer(), numElements * sizeof(half), cudaMemcpyDeviceToHost, stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + saveFp16BinaryFile(path, host.data(), numElements); +} + +//! Select the effective random seed using the configured precedence. +//! +//! Priority is per-request JSON, then top-level JSON, then command line. +//! A JSON seed of zero is treated as "unset", not as a literal seed. +int32_t resolveSeed(int32_t requestSeed, int32_t globalSeed, int32_t cliSeed) +{ + if (requestSeed != 0) + { + return requestSeed; + } + if (globalSeed != 0) + { + return globalSeed; + } + return cliSeed; +} + +//! Select the effective denoising-step count using configured precedence. +//! +//! Priority is per-request JSON, then top-level JSON, then command line. +//! Returning zero tells WFMInferenceRuntime to use its engine configuration. +int32_t resolveNumInferenceSteps(int32_t requestSteps, int32_t globalSteps, int32_t cliSteps) +{ + if (requestSteps > 0) + { + return requestSteps; + } + if (globalSteps > 0) + { + return globalSteps; + } + if (cliSteps > 0) + { + return cliSteps; + } + return 0; +} + +//! Convert one parsed JSON request into the runtime's GPU-backed request type. +//! +//! This is the bridge between configuration/I/O and model execution. It resolves +//! defaults, prepares a pixel tensor, fills dimensions expected by validation, +//! and optionally adds an input waveform. +//! +//! @param spec Per-request values parsed from JSON. +//! @param globals Top-level defaults parsed from JSON. +//! @param args Command-line defaults. +//! @param config Dimensions and sample rate expected by the exported engines. +//! @param stream CUDA stream used while preparing tensors. +//! @return A complete request suitable for WFMInferenceRuntime::handleRequest(). +rt::WFMGenerationRequest buildWfmRequest(WfmRequestSpec const& spec, WfmInputGlobals const& globals, + WfmInferenceArgs const& args, rt::CosmosEngineConfig const& config, cudaStream_t stream) +{ + rt::WFMGenerationRequest request{}; + request.prompt = spec.prompt; + request.generateSound = spec.generateSound; + request.numInferenceSteps + = resolveNumInferenceSteps(spec.numInferenceSteps, globals.numInferenceSteps, args.numInferenceSteps); + request.seed = resolveSeed(spec.seed, globals.seed, args.seed); + + // A supplied pixels_file conditions the request on that tensor. Otherwise, + // random normalized pixels provide a deterministic starting input. + if (!spec.pixelsFile.empty()) + { + request.pixels.buffer = loadPixelsFromFile(spec.pixelsFile, config, stream); + } + else + { + request.pixels.buffer = makeRandomPixels(config, request.seed, stream); + } + request.pixels.batch = 1; + request.pixels.channels = 3; + request.pixels.numFrames = config.numFrames; + request.pixels.height = config.height; + request.pixels.width = config.width; + + // Audio is optional. Leaving inputWaveform.buffer empty tells the runtime + // that this request has no waveform conditioning. + if (!spec.waveformFile.empty()) + { + int64_t numSamples{0}; + request.inputWaveform.buffer = loadWaveformFromFile(spec.waveformFile, config, stream, numSamples); + request.inputWaveform.batch = 1; + request.inputWaveform.sampleRate = config.sampleRate; + request.inputWaveform.numSamples = numSamples; + } + + return request; +} + +//! Parse the input JSON into global defaults and individual request specs. +//! +//! Expected structure: +//! { +//! "num_inference_steps": 2, // optional global default +//! "seed": 42, // optional global default +//! "requests": [ +//! { +//! "prompt": "...", // required +//! "generate_sound": false, // optional +//! "pixels_file": "...", // optional raw FP16 input +//! "waveform_file": "...", // optional raw FP16 input +//! "output_video_file": "...", // optional raw FP16 output +//! "output_waveform_file": "..." // optional raw FP16 output +//! } +//! ] +//! } +//! +//! @param inputFilePath JSON file to parse. +//! @return Pair containing top-level defaults and all request descriptions. +//! @throws std::runtime_error for malformed JSON or invalid required fields. +std::pair> parseInputFile(std::filesystem::path const& inputFilePath) +{ + WfmInputGlobals globals; + std::vector requestSpecs; + + Json inputData; + std::ifstream inputFileStream(inputFilePath); + check::check(inputFileStream.is_open(), "Failed to open input file: " + inputFilePath.string()); + try + { + inputData = Json::parse(inputFileStream); + inputFileStream.close(); + } + catch (Json::parse_error const& e) + { + throw std::runtime_error( + format::fmtstr("Failed to parse input file %s with error: %s", inputFilePath.string().c_str(), e.what())); + } + + globals.numInferenceSteps = inputData.value("num_inference_steps", 0); + globals.seed = inputData.value("seed", 0); + + check::check( + inputData.contains("requests") && inputData["requests"].is_array(), "'requests' array not found in input file"); + + auto const& requestsArray = inputData["requests"]; + for (size_t requestIdx = 0; requestIdx < requestsArray.size(); ++requestIdx) + { + auto const& requestItem = requestsArray[requestIdx]; + check::check(requestItem.is_object(), "Each request must be a JSON object"); + + WfmRequestSpec spec; + // prompt is the only mandatory per-request field. + check::check(requestItem.contains("prompt") && requestItem["prompt"].is_string(), + format::fmtstr("Request %zu must contain a string 'prompt' field", requestIdx)); + spec.prompt = requestItem["prompt"].get(); + check::check(spec.prompt.size() <= limits::tokenizer::kMaxInputTextSizeBytes, + format::fmtstr("Input rejected: prompt too large in request %zu: %zu bytes (max: %zu). Limit defined in " + "%s.", + requestIdx, spec.prompt.size(), limits::tokenizer::kMaxInputTextSizeBytes, + limits::kInputLimitsLocation)); + + spec.generateSound = requestItem.value("generate_sound", false); + spec.numInferenceSteps = requestItem.value("num_inference_steps", 0); + spec.seed = requestItem.value("seed", 0); + spec.pixelsFile = requestItem.value("pixels_file", ""); + spec.waveformFile = requestItem.value("waveform_file", ""); + spec.outputVideoFile = requestItem.value("output_video_file", ""); + spec.outputWaveformFile = requestItem.value("output_waveform_file", ""); + + requestSpecs.push_back(std::move(spec)); + } + + if (requestSpecs.empty()) + { + throw std::runtime_error("No requests found in input file"); + } + + return std::make_pair(std::move(globals), std::move(requestSpecs)); +} + +} // namespace + +//! Program entry point: initialize resources, run all requests, and write results. +//! +//! The requests are processed sequentially on one CUDA stream. This makes tensor +//! lifetime and runtime-owned output buffers straightforward: each response is +//! consumed before the next request begins. +int main(int argc, char* argv[]) +{ + // Add a top-level NVTX range so GPU profiling tools can identify this program. + NVTX_SCOPED_RANGE(nvtx_main, "wfm_inference"); + + // Phase 1: Parse and validate command-line configuration. + WfmInferenceArgs args; + if (!parseWfmInferenceArgs(args, argc, argv)) + { + printUsage(argv[0]); + return EXIT_FAILURE; + } + if (args.help) + { + printUsage(argv[0]); + return EXIT_SUCCESS; + } + + // Profiling is enabled if either console or JSON profiling output was requested. + bool const profilerEnabled = args.dumpProfile || !args.profileOutputFile.empty(); + MemoryMonitor memoryMonitor; + if (profilerEnabled) + { + memoryMonitor.start(); + } + + // TensorRT engines may depend on custom Edge-LLM layers. Loading the plugin + // library before deserializing engines registers those layers with TensorRT. + auto pluginHandles = loadEdgellmPluginLib(); + + // Phase 2: Read lightweight request descriptions from JSON. + WfmInputGlobals globals; + std::vector requestSpecs; + try + { + std::tie(globals, requestSpecs) = parseInputFile(args.inputFile); + LOG_INFO("Successfully parsed %zu requests from input file.", requestSpecs.size()); + } + catch (std::exception const& e) + { + LOG_ERROR("Failed to parse input file: %s", e.what()); + return EXIT_FAILURE; + } + + // Phase 3: Create one CUDA stream. Operations submitted to the same stream + // execute in order, which simplifies synchronization between pipeline stages. + cudaStream_t stream{}; + CUDA_CHECK(cudaStreamCreate(&stream)); + + std::unique_ptr wfmRuntime; + try + { + // Construction loads config.json, packing data, tokenizer data, and the + // TensorRT engines located beneath engineDir. + wfmRuntime = std::make_unique(args.engineDir, stream); + } + catch (std::exception const& e) + { + LOG_ERROR("Failed to initialize WFMInferenceRuntime: %s", e.what()); + CUDA_CHECK(cudaStreamDestroy(stream)); + return EXIT_FAILURE; + } + + auto const& config = wfmRuntime->getEngineConfig(); + + // Phase 4 (optional): Warm up GPU kernels and engine state. Warmup uses the + // first request repeatedly and is deliberately excluded from profiling. + if (args.warmup > 0) + { + setProfilingEnabled(false); + LOG_INFO("Starting warmup with %d runs using the first request...", args.warmup); + auto warmupRequest = buildWfmRequest(requestSpecs[0], globals, args, config, stream); + + for (int32_t warmupRun = 0; warmupRun < args.warmup; ++warmupRun) + { + rt::WFMGenerationResponse warmupResponse; + bool const requestStatus = wfmRuntime->handleRequest(warmupRequest, warmupResponse, stream); + CUDA_CHECK(cudaStreamSynchronize(stream)); + if (!requestStatus) + { + LOG_ERROR("Warmup run %d/%d failed", warmupRun + 1, args.warmup); + CUDA_CHECK(cudaStreamDestroy(stream)); + return EXIT_FAILURE; + } + } + LOG_INFO("Warmup of %d runs completed. Starting actual benchmark runs...", args.warmup); + } + + if (profilerEnabled) + { + setProfilingEnabled(true); + gTimer.reset(); + } + + // Phase 5: Prepare the JSON document that will summarize every response. + Json outputData; + outputData["input_file"] = args.inputFile; + outputData["engine_dir"] = args.engineDir; + outputData["responses"] = Json::array(); + + bool hasFailedRequest = false; + size_t failedCount = 0; + std::string const errorMessage = "TensorRT Edge LLM cannot handle this WFM request."; + + LOG_INFO("Processing %zu requests...", requestSpecs.size()); + for (size_t requestIdx = 0; requestIdx < requestSpecs.size(); ++requestIdx) + { + auto const& spec = requestSpecs[requestIdx]; + rt::WFMGenerationResponse response; + + size_t const progressInterval = std::max(size_t(1), std::min(requestSpecs.size() / 10, size_t(100))); + if ((requestIdx + 1) % progressInterval == 0 || requestIdx == 0 || requestIdx == requestSpecs.size() - 1) + { + LOG_INFO("Progress: %zu/%zu (%f%%)", requestIdx + 1, requestSpecs.size(), + 100.0 * (requestIdx + 1) / requestSpecs.size()); + } + + // Build input GPU tensors only when this request is about to run. + rt::WFMGenerationRequest request = buildWfmRequest(spec, globals, args, config, stream); + + // handleRequest() is the core handoff. Internally the runtime validates + // inputs, prepares text, encodes, denoises, and decodes video/audio. + bool requestStatus = false; + if (profilerEnabled) + { + TIME_STAGE("wfm_inference", stream); + requestStatus = wfmRuntime->handleRequest(request, response, stream); + } + else + { + requestStatus = wfmRuntime->handleRequest(request, response, stream); + } + CUDA_CHECK(cudaStreamSynchronize(stream)); + + // Record settings and success independently of whether tensor files were + // requested. This makes the output JSON useful as a batch manifest. + Json responseJson; + responseJson["request_idx"] = requestIdx; + responseJson["prompt"] = sanitizeUtf8ForJson(spec.prompt); + responseJson["generate_sound"] = spec.generateSound; + responseJson["num_inference_steps"] = request.numInferenceSteps; + responseJson["seed"] = request.seed; + responseJson["success"] = requestStatus; + + if (requestStatus) + { + // --dumpOutput prints tensor shapes only; it does not print tensor data. + if (args.dumpOutput) + { + if (response.outputVideo.buffer) + { + LOG_INFO("Request %zu output video shape: %s", requestIdx, + response.outputVideo.buffer->getShape().formatString().c_str()); + } + if (response.outputWaveform.buffer) + { + LOG_INFO("Request %zu output waveform shape: %s", requestIdx, + response.outputWaveform.buffer->getShape().formatString().c_str()); + } + } + + // If an output path was provided, copy and save the video tensor. + // Otherwise, preserve only its shape in the response JSON. + if (!spec.outputVideoFile.empty() && response.outputVideo.buffer) + { + try + { + saveTensorToFp16File(spec.outputVideoFile, *response.outputVideo.buffer, stream); + responseJson["output_video_file"] = spec.outputVideoFile; + } + catch (std::exception const& e) + { + LOG_ERROR("Failed to write output video for request %zu: %s", requestIdx, e.what()); + responseJson["output_video_file_error"] = e.what(); + hasFailedRequest = true; + ++failedCount; + } + } + else if (response.outputVideo.buffer) + { + responseJson["output_video_shape"] = response.outputVideo.buffer->getShape().formatString(); + } + + // Audio output follows the same policy as video output. It exists only + // when the selected bundle and request actually run the sound pipeline. + if (!spec.outputWaveformFile.empty() && response.outputWaveform.buffer) + { + try + { + saveTensorToFp16File(spec.outputWaveformFile, *response.outputWaveform.buffer, stream); + responseJson["output_waveform_file"] = spec.outputWaveformFile; + } + catch (std::exception const& e) + { + LOG_ERROR("Failed to write output waveform for request %zu: %s", requestIdx, e.what()); + responseJson["output_waveform_file_error"] = e.what(); + hasFailedRequest = true; + ++failedCount; + } + } + else if (response.outputWaveform.buffer) + { + responseJson["output_waveform_shape"] = response.outputWaveform.buffer->getShape().formatString(); + } + } + else + { + hasFailedRequest = true; + ++failedCount; + responseJson["error"] = errorMessage; + LOG_ERROR("*** FAILED *** Request %zu failed to process!", requestIdx); + } + + outputData["responses"].push_back(std::move(responseJson)); + } + + // Phase 6: Stop measurement and report aggregate request status. + LOG_INFO( + "Processing complete: %zu/%zu requests successful", requestSpecs.size() - failedCount, requestSpecs.size()); + if (failedCount > 0) + { + LOG_ERROR("*** %zu REQUESTS FAILED ***", failedCount); + } + + if (profilerEnabled) + { + setProfilingEnabled(false); + memoryMonitor.stop(); + } + + if (args.dumpProfile) + { + std::ostringstream profileOutput; + profileOutput << std::endl; + profileOutput << "=== WFM Performance Summary ===" << std::endl; + outputMemoryProfile(profileOutput, memoryMonitor); + outputLayerProfiles(profileOutput, false); + profileOutput << "=====================================" << std::endl; + LOG_INFO("%s", profileOutput.str().c_str()); + } + + // Profiling JSON is separate from the normal response JSON because it contains + // timing stages and memory measurements rather than model outputs. + if (!args.profileOutputFile.empty()) + { + try + { + Json profileJson; + addJsonTimingStages(profileJson); + addJsonMemorySummary(profileJson, memoryMonitor); + + std::ofstream profileFile(args.profileOutputFile); + if (profileFile.is_open()) + { + profileFile << profileJson.dump(2); + profileFile.close(); + LOG_INFO("Profile data exported to: %s", args.profileOutputFile.c_str()); + } + else + { + LOG_ERROR("Failed to open profile output file: %s", args.profileOutputFile.c_str()); + CUDA_CHECK(cudaStreamDestroy(stream)); + return EXIT_FAILURE; + } + } + catch (std::exception const& e) + { + LOG_ERROR("Failed to write profile output file: %s", e.what()); + CUDA_CHECK(cudaStreamDestroy(stream)); + return EXIT_FAILURE; + } + } + + // Phase 7: Always attempt to write the batch response manifest, including + // entries for requests that failed. + try + { + std::ofstream outputFile(args.outputFile); + if (outputFile.is_open()) + { + outputFile << outputData.dump(4); + outputFile.close(); + LOG_INFO("All responses exported to: %s", args.outputFile.c_str()); + } + else + { + LOG_ERROR("Failed to open output file: %s", args.outputFile.c_str()); + CUDA_CHECK(cudaStreamDestroy(stream)); + return EXIT_FAILURE; + } + } + catch (std::exception const& e) + { + LOG_ERROR("Failed to write output file: %s", e.what()); + CUDA_CHECK(cudaStreamDestroy(stream)); + return EXIT_FAILURE; + } + + // Release the CUDA stream after all asynchronous work and output copies finish. + CUDA_CHECK(cudaStreamDestroy(stream)); + // A partial batch failure produces a failing process exit code even though + // successful request entries are still present in the output JSON. + return hasFailedRequest ? EXIT_FAILURE : EXIT_SUCCESS; +} diff --git a/examples/wfm/wfm_input_example.json b/examples/wfm/wfm_input_example.json new file mode 100644 index 00000000..231a6689 --- /dev/null +++ b/examples/wfm/wfm_input_example.json @@ -0,0 +1,11 @@ +{ + "num_inference_steps": 2, + "seed": 42, + "requests": [ + { + "prompt": "A robot arm picks up a red cube.", + "generate_sound": false, + "output_video_file": "/tmp/wfm_output_video.fp16" + } + ] +} diff --git a/tensorrt_edgellm/models/default/modeling_default.py b/tensorrt_edgellm/models/default/modeling_default.py index e87573e1..b8deacf3 100644 --- a/tensorrt_edgellm/models/default/modeling_default.py +++ b/tensorrt_edgellm/models/default/modeling_default.py @@ -518,6 +518,9 @@ class CausalLM(nn.Module): #: Subclasses override to True when the model must emit ``hidden_states`` #: as an ONNX output in addition to ``logits``. emit_hidden_states: bool = False + #: Debug/export option for emitting the final RMSNorm output. Unlike + #: ``emit_hidden_states``, this matches HuggingFace ``last_hidden_state``. + emit_normed_hidden_states: bool = False def __init__(self, config: ModelConfig) -> None: super().__init__() @@ -632,7 +635,9 @@ def onnx_export_spec(self) -> OnnxSpec: ] + [f"deepstack_embeds_{i}" for i in range(Nd)]) output_names = (["logits"] + [f"present_key_values_{i}" for i in range(Na)]) - if self.emit_hidden_states and not eagle_base: + emit_hidden_output = (self.emit_hidden_states + or self.emit_normed_hidden_states) + if emit_hidden_output and not eagle_base: output_names = (["logits", "hidden_states"] + [f"present_key_values_{i}" for i in range(Na)]) @@ -688,12 +693,11 @@ def onnx_export_spec(self) -> OnnxSpec: 2: mask_kv_len }) # attention_mask - wrapped = _make_flat_wrapper( - self, - Na, - Nd, - eagle_base=eagle_base, - emit_hidden_states=self.emit_hidden_states) + wrapped = _make_flat_wrapper(self, + Na, + Nd, + eagle_base=eagle_base, + emit_hidden_states=emit_hidden_output) wrapped.eval() return OnnxSpec(wrapped=wrapped, @@ -766,6 +770,9 @@ def forward( dim=-1).to(torch.float16) return logits, eagle_hidden, present_key_values + if self.emit_normed_hidden_states: + return logits, hidden_states, present_key_values + if self.emit_hidden_states: # Full-sequence last-layer pre-norm residual, populated by # :meth:`Transformer.forward` (see its docstring for the HF diff --git a/tensorrt_edgellm/scripts/export.py b/tensorrt_edgellm/scripts/export.py index fa279fcd..019ac317 100644 --- a/tensorrt_edgellm/scripts/export.py +++ b/tensorrt_edgellm/scripts/export.py @@ -511,6 +511,7 @@ def _export_llm(model_dir: str, llm_out_dir: str, model_type: str = "", eagle_base: bool = False, + emit_normed_hidden_states: bool = False, fp8_embedding: bool = False, reduced_vocab_dir: str = "", mtp_base: bool = False, @@ -573,6 +574,8 @@ def _export_llm(model_dir: str, tp_size=world, tp_rank=rank, ) + if emit_normed_hidden_states: + model.emit_normed_hidden_states = True except (OSError, ValueError, RuntimeError, ImportError) as exc: logger.exception("[LLM] Failed to load checkpoint") raise SystemExit(1) from exc @@ -1829,6 +1832,13 @@ def main() -> None: help= "Write embedding.safetensors in FP8 E4M3 format with per-row block scales.", ) + p.add_argument( + "--emit-normed-hidden-states", + action="store_true", + help=( + "Add the full-sequence final RMSNorm hidden_states output to the " + "LLM ONNX graph for eager parity debugging."), + ) p.add_argument( "--reduced-vocab-dir", "--reduced_vocab_dir", @@ -2009,18 +2019,19 @@ def _allow(component: str) -> bool: # drive both the pre-run log and the post-run summary below. stages = [ (_has_llm_component(model_type, "thinker") and not args.skip_llm - and not _draft_only and _allow("thinker"), "thinker", - lambda out: _export_llm(model_dir, - out, - model_type=model_type, - eagle_base=args.eagle_base, - mtp_base=args.mtp, - dflash_base=args.dflash_base, - dflash_draft_dir=args.dflash_draft_dir, - fp8_embedding=args.fp8_embedding, - reduced_vocab_dir=args.reduced_vocab_dir, - externalize_weights=externalize_weights, - tp_size=args.tp_size)), + and not _draft_only and _allow("thinker"), "thinker", lambda out: + _export_llm(model_dir, + out, + model_type=model_type, + eagle_base=args.eagle_base, + emit_normed_hidden_states=args.emit_normed_hidden_states, + mtp_base=args.mtp, + dflash_base=args.dflash_base, + dflash_draft_dir=args.dflash_draft_dir, + fp8_embedding=args.fp8_embedding, + reduced_vocab_dir=args.reduced_vocab_dir, + externalize_weights=externalize_weights, + tp_size=args.tp_size)), (args.mtp, "mtp_draft", lambda out: _export_mtp_draft( model_dir, out, externalize_weights=externalize_weights)), (args.dflash_draft, "dflash_draft", lambda out: _export_dflash_draft(