From 8926ac3d8a64656bb7283ba2ef6f6c9c677d07aa Mon Sep 17 00:00:00 2001 From: Michael Williams Date: Tue, 22 Sep 2026 14:33:49 -0700 Subject: [PATCH] Add vision-only Edge-LLM ExecuTorch delegate integration --- cpp/BUILD | 23 +- .../executorch/EdgeLLMBackend.h | 42 ++++ .../executorch/EdgeLLMBlobHeader.h | 26 ++ .../torch_tensorrt/executorch/CMakeLists.txt | 52 ++++ .../executorch/EdgeLLMBackend.cpp | 229 ++++++++++++++++++ .../executorch/EdgeLLMBlobHeader.cpp | 177 ++++++++++++++ .../CMakeLists.txt | 6 + examples/executorch_reference_runner/main.cpp | 38 +-- py/torch-tensorrt-executorch-runtime/setup.py | 8 +- tests/cpp/executorch/BUILD | 10 + .../executorch/test_edge_llm_blob_header.cpp | 103 ++++++++ .../executorch/test_edge_llm_partitioner.py | 205 ++++++++++++++++ tools/hf/export_pi05_vision_executorch.py | 78 ++++++ tools/hf/exporters/executorch/__init__.py | 55 +++++ tools/hf/exporters/executorch/artifact.py | 101 ++++++++ tools/hf/exporters/executorch/backend.py | 103 ++++++++ .../exporters/executorch/operator_support.py | 22 ++ tools/hf/exporters/executorch/partitioner.py | 73 ++++++ .../hf/exporters/executorch/serialization.py | 138 +++++++++++ tools/hf/exporters/executorch/vision.py | 65 +++++ tools/hf/exporters/models/pi05/spec.py | 16 +- tools/hf/exporters/models/pi05/vision.py | 36 +++ tools/hf/exporters/ops.py | 160 +++++++++++- .../hf/exporters/tests/test_edge_exporter.py | 65 ++++- 24 files changed, 1798 insertions(+), 33 deletions(-) create mode 100644 cpp/include/torch_tensorrt/executorch/EdgeLLMBackend.h create mode 100644 cpp/include/torch_tensorrt/executorch/EdgeLLMBlobHeader.h create mode 100644 cpp/src/torch_tensorrt/executorch/EdgeLLMBackend.cpp create mode 100644 cpp/src/torch_tensorrt/executorch/EdgeLLMBlobHeader.cpp create mode 100644 tests/cpp/executorch/test_edge_llm_blob_header.cpp create mode 100644 tests/py/dynamo/executorch/test_edge_llm_partitioner.py create mode 100644 tools/hf/export_pi05_vision_executorch.py create mode 100644 tools/hf/exporters/executorch/__init__.py create mode 100644 tools/hf/exporters/executorch/artifact.py create mode 100644 tools/hf/exporters/executorch/backend.py create mode 100644 tools/hf/exporters/executorch/operator_support.py create mode 100644 tools/hf/exporters/executorch/partitioner.py create mode 100644 tools/hf/exporters/executorch/serialization.py create mode 100644 tools/hf/exporters/executorch/vision.py create mode 100644 tools/hf/exporters/models/pi05/vision.py diff --git a/cpp/BUILD b/cpp/BUILD index 30619cda923..e7c483c3b35 100644 --- a/cpp/BUILD +++ b/cpp/BUILD @@ -119,6 +119,17 @@ cc_library( strip_include_prefix = "include", ) +cc_library( + name = "edge_llm_executorch_blob_header", + srcs = [ + "src/torch_tensorrt/executorch/EdgeLLMBlobHeader.cpp", + ], + hdrs = [ + "include/torch_tensorrt/executorch/EdgeLLMBlobHeader.h", + ], + strip_include_prefix = "include", +) + cc_library( name = "tensorrt_executorch_weight_streaming_budget", srcs = [ @@ -153,7 +164,6 @@ cc_library( srcs = [ "src/torch_tensorrt/executorch/RegisterCudaDeviceAllocator.cpp", ], - alwayslink = True, target_compatible_with = select({ ":linux_x86_64": [], ":sbsa": [], @@ -170,6 +180,7 @@ cc_library( ], "//conditions:default": [], }), + alwayslink = True, ) # Registration only, for the same reason as the allocator above: the kernel @@ -184,11 +195,11 @@ cc_library( srcs = [ "src/torch_tensorrt/executorch/RegisterDeviceCopyKernels.cpp", ], - alwayslink = True, deps = [ "@executorch//:executorch_device_copy_kernels", "@executorch//:executorch_headers", ], + alwayslink = True, ) cc_library( @@ -214,15 +225,15 @@ cc_library( ":tensorrt_executorch_weight_streaming_budget", ] + select({ ":linux_x86_64": [ + "@cuda//:cudart", "@executorch//:executorch_headers", "@executorch//:extension_cuda", - "@cuda//:cudart", "@tensorrt//:nvinfer", ], ":sbsa": [ + "@cuda//:cudart", "@executorch//:executorch_headers", "@executorch//:extension_cuda", - "@cuda//:cudart", "@tensorrt_sbsa//:nvinfer", ], "//conditions:default": [], @@ -234,6 +245,8 @@ filegroup( name = "executorch_backend_source_files", srcs = [ "src/torch_tensorrt/executorch/CMakeLists.txt", + "src/torch_tensorrt/executorch/EdgeLLMBackend.cpp", + "src/torch_tensorrt/executorch/EdgeLLMBlobHeader.cpp", "src/torch_tensorrt/executorch/README.md", "src/torch_tensorrt/executorch/TensorRTBackend.cpp", "src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp", @@ -254,6 +267,8 @@ filegroup( filegroup( name = "executorch_api_headers", srcs = [ + "include/torch_tensorrt/executorch/EdgeLLMBackend.h", + "include/torch_tensorrt/executorch/EdgeLLMBlobHeader.h", "include/torch_tensorrt/executorch/TensorRTBackend.h", "include/torch_tensorrt/executorch/TensorRTBindingNames.h", "include/torch_tensorrt/executorch/TensorRTBlobHeader.h", diff --git a/cpp/include/torch_tensorrt/executorch/EdgeLLMBackend.h b/cpp/include/torch_tensorrt/executorch/EdgeLLMBackend.h new file mode 100644 index 00000000000..684769562ad --- /dev/null +++ b/cpp/include/torch_tensorrt/executorch/EdgeLLMBackend.h @@ -0,0 +1,42 @@ +#pragma once + +#include +#include + +#include + +#include +#include + +namespace torch_tensorrt { +namespace executorch_backend { + +struct EdgeLLMHandle { + int device_id = 0; + std::unique_ptr vision_runner; + std::mutex mu; + cudaEvent_t inflight_event = nullptr; + bool inflight_pending = false; + + ~EdgeLLMHandle(); +}; + +class EdgeLLMBackend final : public ::executorch::runtime::BackendInterface { + public: + bool is_available() const override; + + ::executorch::runtime::Result<::executorch::runtime::DelegateHandle*> init( + ::executorch::runtime::BackendInitContext& context, + ::executorch::runtime::FreeableBuffer* processed, + ::executorch::runtime::ArrayRef<::executorch::runtime::CompileSpec> compile_specs) const override; + + ::executorch::runtime::Error execute( + ::executorch::runtime::BackendExecutionContext& context, + ::executorch::runtime::DelegateHandle* handle, + ::executorch::runtime::Span<::executorch::runtime::EValue*> args) const override; + + void destroy(::executorch::runtime::DelegateHandle* handle) const override; +}; + +} // namespace executorch_backend +} // namespace torch_tensorrt diff --git a/cpp/include/torch_tensorrt/executorch/EdgeLLMBlobHeader.h b/cpp/include/torch_tensorrt/executorch/EdgeLLMBlobHeader.h new file mode 100644 index 00000000000..64b8d4f5f42 --- /dev/null +++ b/cpp/include/torch_tensorrt/executorch/EdgeLLMBlobHeader.h @@ -0,0 +1,26 @@ +#pragma once + +#include +#include +#include + +namespace torch_tensorrt { +namespace executorch_backend { + +struct EdgeLLMBlobHeader { + uint32_t metadata_offset = 0; + uint32_t metadata_size = 0; + uint32_t blob_offset = 0; + uint64_t blob_size = 0; + int abi_version = 0; + std::string component; + std::string runner; + std::string metadata_json; + std::string runner_config_json; + + static const void* nested_blob_data(const void* payload, const EdgeLLMBlobHeader& h); + static bool parse(const void* data, std::size_t size, EdgeLLMBlobHeader& out); +}; + +} // namespace executorch_backend +} // namespace torch_tensorrt diff --git a/cpp/src/torch_tensorrt/executorch/CMakeLists.txt b/cpp/src/torch_tensorrt/executorch/CMakeLists.txt index 1b503c567da..bc70dde088c 100644 --- a/cpp/src/torch_tensorrt/executorch/CMakeLists.txt +++ b/cpp/src/torch_tensorrt/executorch/CMakeLists.txt @@ -23,6 +23,7 @@ find_package(Threads REQUIRED) set(_torchtrt_executorch_sources "${CMAKE_CURRENT_LIST_DIR}/TensorRTBackend.cpp" "${CMAKE_CURRENT_LIST_DIR}/TensorRTBlobHeader.cpp" + "${CMAKE_CURRENT_LIST_DIR}/EdgeLLMBlobHeader.cpp" "${CMAKE_CURRENT_LIST_DIR}/WeightStreamingBudget.cpp" ) @@ -189,6 +190,57 @@ else() ) endif() +option( + TORCHTRT_BUILD_EDGE_LLM_EXECUTORCH_BACKEND + "Build the optional TensorRT-Edge-LLM ExecuTorch backend" + OFF +) +if(TORCHTRT_BUILD_EDGE_LLM_EXECUTORCH_BACKEND) + set(EDGELLM_INCLUDE_DIR "" CACHE PATH "TensorRT-Edge-LLM cpp include root") + set(EDGELLM_EXECUTORCH_LIBRARY "" CACHE FILEPATH "Path to libedgellmExecutorch.so") + foreach(_required IN ITEMS EDGELLM_INCLUDE_DIR EDGELLM_EXECUTORCH_LIBRARY) + if("${${_required}}" STREQUAL "" OR NOT EXISTS "${${_required}}") + message(FATAL_ERROR "${_required} must name an existing TensorRT-Edge-LLM artifact") + endif() + endforeach() + + add_library(edgellm_executorch SHARED IMPORTED) + set_target_properties(edgellm_executorch PROPERTIES IMPORTED_LOCATION "${EDGELLM_EXECUTORCH_LIBRARY}") + + add_library(executorch_edge_llm_backend STATIC + "${CMAKE_CURRENT_LIST_DIR}/EdgeLLMBackend.cpp" + ) + target_include_directories(executorch_edge_llm_backend + PUBLIC + "${CMAKE_CURRENT_LIST_DIR}/../../../include" + "${EDGELLM_INCLUDE_DIR}" + ) + target_compile_definitions(executorch_edge_llm_backend + PUBLIC + C10_USING_CUSTOM_GENERATED_MACROS + ) + target_link_libraries(executorch_edge_llm_backend + PUBLIC + executorch_trt_backend + edgellm_executorch + ${_torchtrt_executorch_link_libraries} + ${CMAKE_DL_LIBS} + ) + + add_library(torchtrt_edge_llm_executorch_backend INTERFACE) + add_library(torchtrt::edge_llm_executorch_backend ALIAS torchtrt_edge_llm_executorch_backend) + if(MSVC) + target_link_libraries(torchtrt_edge_llm_executorch_backend + INTERFACE executorch_edge_llm_backend edgellm_executorch) + else() + target_link_libraries(torchtrt_edge_llm_executorch_backend + INTERFACE + "-Wl,--whole-archive,$,--no-whole-archive" + edgellm_executorch + ) + endif() +endif() + install( TARGETS executorch_trt_backend ARCHIVE DESTINATION "${CMAKE_INSTALL_LIBDIR}" diff --git a/cpp/src/torch_tensorrt/executorch/EdgeLLMBackend.cpp b/cpp/src/torch_tensorrt/executorch/EdgeLLMBackend.cpp new file mode 100644 index 00000000000..fabae21c5bb --- /dev/null +++ b/cpp/src/torch_tensorrt/executorch/EdgeLLMBackend.cpp @@ -0,0 +1,229 @@ +#include "torch_tensorrt/executorch/EdgeLLMBackend.h" +#include "torch_tensorrt/executorch/EdgeLLMBlobHeader.h" +#include "torch_tensorrt/executorch/TensorRTBlobHeader.h" + +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace torch_tensorrt { +namespace executorch_backend { + +using ::executorch::runtime::ArrayRef; +using ::executorch::runtime::BackendExecutionContext; +using ::executorch::runtime::BackendInitContext; +using ::executorch::runtime::CompileSpec; +using ::executorch::runtime::DelegateHandle; +using ::executorch::runtime::Error; +using ::executorch::runtime::EValue; +using ::executorch::runtime::FreeableBuffer; +using ::executorch::runtime::MemoryAllocator; +using ::executorch::runtime::Result; +using ::executorch::runtime::Span; + +namespace { + +extern const Error kEdgeLLMRegistrationResult; + +Error check_registration() { + if (kEdgeLLMRegistrationResult != Error::Ok) { + ET_LOG( + Error, "EdgeLLMBackend registration failed: %s", ::executorch::runtime::to_string(kEdgeLLMRegistrationResult)); + } + return kEdgeLLMRegistrationResult; +} + +bool is_cuda_accessible(const void* pointer) { + if (pointer == nullptr) { + return false; + } + cudaPointerAttributes attributes{}; + const cudaError_t status = cudaPointerGetAttributes(&attributes, pointer); + if (status != cudaSuccess) { + cudaGetLastError(); + return false; + } + return attributes.type == cudaMemoryTypeDevice || attributes.type == cudaMemoryTypeManaged; +} + +std::vector tensor_shape(const ::executorch::aten::Tensor& tensor) { + std::vector shape; + shape.reserve(static_cast(tensor.dim())); + for (ssize_t dim = 0; dim < tensor.dim(); ++dim) { + shape.push_back(static_cast(tensor.size(dim))); + } + return shape; +} + +struct HandleDeleter { + void operator()(EdgeLLMHandle* handle) const { + if (handle != nullptr) { + handle->~EdgeLLMHandle(); + } + } +}; + +} // namespace + +EdgeLLMHandle::~EdgeLLMHandle() { + int entry_device = -1; + const bool restore_device = cudaGetDevice(&entry_device) == cudaSuccess && entry_device != device_id; + (void)cudaSetDevice(device_id); + if (inflight_event != nullptr && inflight_pending) { + (void)cudaEventSynchronize(inflight_event); + inflight_pending = false; + } + vision_runner.reset(); + if (inflight_event != nullptr) { + (void)cudaEventDestroy(inflight_event); + inflight_event = nullptr; + } + if (restore_device) { + (void)cudaSetDevice(entry_device); + } +} + +bool EdgeLLMBackend::is_available() const { + return check_registration() == Error::Ok; +} + +Result EdgeLLMBackend::init( + BackendInitContext& context, + FreeableBuffer* processed, + ArrayRef compile_specs) const { + (void)compile_specs; + if (check_registration() != Error::Ok) { + return kEdgeLLMRegistrationResult; + } + if (processed == nullptr || processed->data() == nullptr) { + ET_LOG(Error, "EdgeLLMBackend::init: null processed payload"); + return Error::InvalidArgument; + } + + EdgeLLMBlobHeader edge_header; + if (!EdgeLLMBlobHeader::parse(processed->data(), processed->size(), edge_header)) { + ET_LOG(Error, "EdgeLLMBackend::init: invalid EL01 component payload"); + return Error::InvalidProgram; + } + const void* nested_blob = EdgeLLMBlobHeader::nested_blob_data(processed->data(), edge_header); + TensorRTBlobHeader trt_header; + if (!TensorRTBlobHeader::parse(nested_blob, static_cast(edge_header.blob_size), trt_header)) { + ET_LOG(Error, "EdgeLLMBackend::init: invalid nested TensorRT payload"); + return Error::InvalidProgram; + } + if (trt_header.input_binding_names.size() != 1 || trt_header.output_binding_names.size() != 1) { + ET_LOG(Error, "EdgeLLMBackend::init: vision runner requires exactly one input and one output"); + return Error::InvalidProgram; + } + + MemoryAllocator* allocator = context.get_runtime_allocator(); + if (allocator == nullptr) { + return Error::InvalidState; + } + EdgeLLMHandle* handle = allocator->allocateInstance(); + if (handle == nullptr) { + return Error::MemoryAllocationFailed; + } + new (handle) EdgeLLMHandle(); + std::unique_ptr handle_guard(handle); + handle->device_id = trt_header.device_id; + + if (cudaSetDevice(handle->device_id) != cudaSuccess) { + ET_LOG(Error, "EdgeLLMBackend::init: failed to select CUDA device %d", handle->device_id); + return Error::InvalidProgram; + } + if (cudaEventCreateWithFlags(&handle->inflight_event, cudaEventDisableTiming | cudaEventBlockingSync) != + cudaSuccess) { + ET_LOG(Error, "EdgeLLMBackend::init: failed to create completion event"); + return Error::InvalidProgram; + } + + const auto caller_stream = ::executorch::extension::cuda::getCallerStream(); + cudaStream_t stream = caller_stream.value_or(cudaStreamPerThread); + const void* engine_data = TensorRTBlobHeader::engine_data(nested_blob, trt_header); + handle->vision_runner = trt_edgellm::executorch::VitExecutorchAdapter::create( + {engine_data, static_cast(trt_header.engine_size)}, stream); + if (!handle->vision_runner) { + ET_LOG(Error, "EdgeLLMBackend::init: failed to create VitRunner adapter"); + return Error::InvalidProgram; + } + + processed->Free(); + handle_guard.release(); + return static_cast(handle); +} + +Error EdgeLLMBackend::execute(BackendExecutionContext& context, DelegateHandle* delegate_handle, Span args) + const { + (void)context; + if (delegate_handle == nullptr || args.size() != 2 || args[0] == nullptr || args[1] == nullptr || + !args[0]->isTensor() || !args[1]->isTensor()) { + ET_LOG(Error, "EdgeLLMBackend::execute: expected one tensor input and one tensor output"); + return Error::InvalidArgument; + } + auto* handle = static_cast(delegate_handle); + std::lock_guard lock(handle->mu); + if (handle->inflight_pending) { + if (cudaEventSynchronize(handle->inflight_event) != cudaSuccess) { + return Error::InvalidProgram; + } + handle->inflight_pending = false; + } + + auto input = args[0]->toTensor(); + auto output = args[1]->toTensor(); + if (input.scalar_type() != ::executorch::aten::ScalarType::Half || + output.scalar_type() != ::executorch::aten::ScalarType::Half || !is_cuda_accessible(input.const_data_ptr()) || + !is_cuda_accessible(output.mutable_data_ptr())) { + ET_LOG(Error, "EdgeLLMBackend::execute: prepared vision input/output must be device-resident FP16 tensors"); + return Error::InvalidArgument; + } + + const auto caller_stream = ::executorch::extension::cuda::getCallerStream(); + cudaStream_t stream = caller_stream.value_or(cudaStreamPerThread); + const bool ok = handle->vision_runner->execute( + {input.mutable_data_ptr(), tensor_shape(input), nvinfer1::DataType::kHALF}, + {output.mutable_data_ptr(), tensor_shape(output), nvinfer1::DataType::kHALF}, + stream); + if (!ok) { + return Error::InvalidProgram; + } + + if (cudaEventRecord(handle->inflight_event, stream) != cudaSuccess) { + (void)cudaStreamSynchronize(stream); + return Error::InvalidProgram; + } + handle->inflight_pending = true; + return Error::Ok; +} + +void EdgeLLMBackend::destroy(DelegateHandle* handle) const { + if (handle != nullptr) { + static_cast(handle)->~EdgeLLMHandle(); + } +} + +} // namespace executorch_backend +} // namespace torch_tensorrt + +namespace torch_tensorrt { +namespace executorch_backend { +namespace { + +EdgeLLMBackend& get_edge_llm_backend() { + static EdgeLLMBackend backend; + return backend; +} + +const ::executorch::runtime::Backend kEdgeLLMBackendId{"EdgeLLMBackend", &get_edge_llm_backend()}; +const Error kEdgeLLMRegistrationResult = ::executorch::runtime::register_backend(kEdgeLLMBackendId); + +} // namespace +} // namespace executorch_backend +} // namespace torch_tensorrt diff --git a/cpp/src/torch_tensorrt/executorch/EdgeLLMBlobHeader.cpp b/cpp/src/torch_tensorrt/executorch/EdgeLLMBlobHeader.cpp new file mode 100644 index 00000000000..8b634c63427 --- /dev/null +++ b/cpp/src/torch_tensorrt/executorch/EdgeLLMBlobHeader.cpp @@ -0,0 +1,177 @@ +#include "torch_tensorrt/executorch/EdgeLLMBlobHeader.h" + +#include +#include +#include +#include +#include + +namespace torch_tensorrt { +namespace executorch_backend { +namespace { + +constexpr char EDGE_LLM_MAGIC[4] = {'E', 'L', '0', '1'}; +constexpr uint32_t METADATA_OFFSET_FIELD_OFFSET = 4; +constexpr uint32_t METADATA_SIZE_FIELD_OFFSET = 8; +constexpr uint32_t BLOB_OFFSET_FIELD_OFFSET = 12; +constexpr uint32_t BLOB_SIZE_FIELD_OFFSET = 16; +constexpr uint32_t HEADER_SIZE = 32; +constexpr uint32_t BLOB_ALIGNMENT = 16; +constexpr int SUPPORTED_ABI_VERSION = 1; + +template +T read_field(const uint8_t* data, std::size_t offset) { + T value{}; + std::memcpy(&value, data + offset, sizeof(T)); + return value; +} + +std::size_t skip_ws(const std::string& value, std::size_t pos) { + while (pos < value.size() && (value[pos] == ' ' || value[pos] == '\t' || value[pos] == '\n' || value[pos] == '\r')) { + ++pos; + } + return pos; +} + +std::size_t value_after_key(const std::string& json, const char* key) { + const std::string quoted_key = std::string("\"") + key + "\""; + const std::size_t key_pos = json.find(quoted_key); + if (key_pos == std::string::npos) { + return std::string::npos; + } + const std::size_t colon = json.find(':', key_pos + quoted_key.size()); + if (colon == std::string::npos) { + return std::string::npos; + } + return skip_ws(json, colon + 1); +} + +bool parse_int(const std::string& json, const char* key, int& out) { + std::size_t pos = value_after_key(json, key); + if (pos == std::string::npos || pos >= json.size()) { + return false; + } + bool negative = false; + if (json[pos] == '-') { + negative = true; + ++pos; + } + int value = 0; + bool saw_digit = false; + while (pos < json.size() && json[pos] >= '0' && json[pos] <= '9') { + saw_digit = true; + if (value > (std::numeric_limits::max() - (json[pos] - '0')) / 10) { + return false; + } + value = value * 10 + (json[pos] - '0'); + ++pos; + } + if (!saw_digit) { + return false; + } + out = negative ? -value : value; + return true; +} + +bool parse_string(const std::string& json, const char* key, std::string& out) { + std::size_t pos = value_after_key(json, key); + if (pos == std::string::npos || pos >= json.size() || json[pos] != '"') { + return false; + } + ++pos; + out.clear(); + while (pos < json.size()) { + if (json[pos] == '"') { + return true; + } + if (json[pos] == '\\') { + ++pos; + if (pos >= json.size()) { + return false; + } + } + out.push_back(json[pos++]); + } + return false; +} + +bool parse_compound(const std::string& json, const char* key, char open, char close, std::string& out) { + std::size_t pos = value_after_key(json, key); + if (pos == std::string::npos || pos >= json.size() || json[pos] != open) { + return false; + } + const std::size_t start = pos; + int depth = 0; + bool in_string = false; + bool escaped = false; + for (; pos < json.size(); ++pos) { + const char ch = json[pos]; + if (in_string) { + if (escaped) { + escaped = false; + } else if (ch == '\\') { + escaped = true; + } else if (ch == '"') { + in_string = false; + } + continue; + } + if (ch == '"') { + in_string = true; + } else if (ch == open) { + ++depth; + } else if (ch == close) { + --depth; + if (depth == 0) { + out = json.substr(start, pos - start + 1); + return true; + } + } + } + return false; +} + +bool parse_metadata(const std::string& json, EdgeLLMBlobHeader& out) { + std::string outputs_json; + return parse_int(json, "abi_version", out.abi_version) && out.abi_version == SUPPORTED_ABI_VERSION && + parse_string(json, "component", out.component) && out.component == "vision" && + parse_string(json, "runner", out.runner) && out.runner == "vit" && + parse_compound(json, "outputs", '[', ']', outputs_json) && outputs_json != "[]" && + parse_compound(json, "runner_config", '{', '}', out.runner_config_json); +} + +} // namespace + +const void* EdgeLLMBlobHeader::nested_blob_data(const void* payload, const EdgeLLMBlobHeader& h) { + return static_cast(payload) + h.blob_offset; +} + +bool EdgeLLMBlobHeader::parse(const void* data, std::size_t size, EdgeLLMBlobHeader& out) { + if (data == nullptr || size < HEADER_SIZE) { + return false; + } + const auto* bytes = static_cast(data); + if (std::memcmp(bytes, EDGE_LLM_MAGIC, sizeof(EDGE_LLM_MAGIC)) != 0) { + return false; + } + + out = EdgeLLMBlobHeader{}; + out.metadata_offset = read_field(bytes, METADATA_OFFSET_FIELD_OFFSET); + out.metadata_size = read_field(bytes, METADATA_SIZE_FIELD_OFFSET); + out.blob_offset = read_field(bytes, BLOB_OFFSET_FIELD_OFFSET); + out.blob_size = read_field(bytes, BLOB_SIZE_FIELD_OFFSET); + + const uint64_t metadata_end = static_cast(out.metadata_offset) + out.metadata_size; + const uint64_t blob_end = static_cast(out.blob_offset) + out.blob_size; + if (out.metadata_offset < HEADER_SIZE || out.blob_offset % BLOB_ALIGNMENT != 0 || metadata_end > out.blob_offset || + blob_end > size) { + return false; + } + + out.metadata_json.assign( + reinterpret_cast(bytes + out.metadata_offset), static_cast(out.metadata_size)); + return parse_metadata(out.metadata_json, out); +} + +} // namespace executorch_backend +} // namespace torch_tensorrt diff --git a/examples/executorch_reference_runner/CMakeLists.txt b/examples/executorch_reference_runner/CMakeLists.txt index 3c30ef81af5..22c9e4a586e 100644 --- a/examples/executorch_reference_runner/CMakeLists.txt +++ b/examples/executorch_reference_runner/CMakeLists.txt @@ -66,6 +66,12 @@ target_link_libraries( executorch::extensions executorch::kernels torchtrt::executorch_backend) +if(TORCHTRT_BUILD_EDGE_LLM_EXECUTORCH_BACKEND) + target_link_libraries( + example_executorch_runner + PRIVATE + torchtrt::edge_llm_executorch_backend) +endif() # Caller-owned KV-cache persistence check (see kv_cache_decode_check.cpp). It # cudaMalloc's the device-tagged planned arenas that hold the KV buffers and diff --git a/examples/executorch_reference_runner/main.cpp b/examples/executorch_reference_runner/main.cpp index a5f4d89594d..b611fc65f38 100644 --- a/examples/executorch_reference_runner/main.cpp +++ b/examples/executorch_reference_runner/main.cpp @@ -55,8 +55,8 @@ using executorch::runtime::MethodMeta; using executorch::runtime::Program; using executorch::runtime::Result; using executorch::runtime::Span; -using executorch::runtime::etensor::Device; using executorch::runtime::TensorInfo; +using executorch::runtime::etensor::Device; static uint8_t method_allocator_pool[4 * 1024U * 1024U]; static uint8_t temp_allocator_pool[1 * 1024U * 1024U]; @@ -71,7 +71,6 @@ static const char* get_flag(int argc, char** argv, const char* flag, const char* return def; } - // The CUDA driver API is resolved at runtime rather than linked. The release build // image ships neither libcuda nor a stub, so linking it would break the build for // everyone to serve one optional flag, and it would have to be wired into both this @@ -84,7 +83,12 @@ struct CudaDriverApi { CUresult (*DeviceGet)(CUdevice*, int) = nullptr; CUresult (*DeviceGetDevResource)(CUdevice, CUdevResource*, CUdevResourceType) = nullptr; CUresult (*DevSmResourceSplitByCount)( - CUdevResource*, unsigned int*, const CUdevResource*, CUdevResource*, unsigned int, unsigned int) = nullptr; + CUdevResource*, + unsigned int*, + const CUdevResource*, + CUdevResource*, + unsigned int, + unsigned int) = nullptr; CUresult (*DevResourceGenerateDesc)(CUdevResourceDesc*, CUdevResource*, unsigned int) = nullptr; CUresult (*GreenCtxCreate)(CUgreenCtx*, CUdevResourceDesc, CUdevice, unsigned int) = nullptr; CUresult (*GreenCtxStreamCreate)(CUstream*, CUgreenCtx, unsigned int, int) = nullptr; @@ -119,10 +123,8 @@ const CudaDriverApi* load_cuda_driver_api() { const bool ok = bind(api.Init, "cuInit") && bind(api.DeviceGet, "cuDeviceGet") && bind(api.DeviceGetDevResource, "cuDeviceGetDevResource") && bind(api.DevSmResourceSplitByCount, "cuDevSmResourceSplitByCount") && - bind(api.DevResourceGenerateDesc, "cuDevResourceGenerateDesc") && - bind(api.GreenCtxCreate, "cuGreenCtxCreate") && - bind(api.GreenCtxStreamCreate, "cuGreenCtxStreamCreate") && - bind(api.GreenCtxDestroy, "cuGreenCtxDestroy") && + bind(api.DevResourceGenerateDesc, "cuDevResourceGenerateDesc") && bind(api.GreenCtxCreate, "cuGreenCtxCreate") && + bind(api.GreenCtxStreamCreate, "cuGreenCtxStreamCreate") && bind(api.GreenCtxDestroy, "cuGreenCtxDestroy") && bind(api.GetErrorString, "cuGetErrorString"); loaded = ok; return ok ? &api : nullptr; @@ -281,11 +283,7 @@ int main(int argc, char** argv) { static_cast(buffer_device->type()), static_cast(device_buffer.error())); ET_LOG( - Info, - " planned buffer[%zu] = %zu bytes on device_type %d", - i, - sz, - static_cast(buffer_device->type())); + Info, " planned buffer[%zu] = %zu bytes on device_type %d", i, sz, static_cast(buffer_device->type())); planned_spans.push_back(device_buffer->as_span()); planned_device_buffers.push_back(std::move(device_buffer.get())); } @@ -299,7 +297,7 @@ int main(int argc, char** argv) { ET_LOG(Info, "Method loaded. inputs=%zu outputs=%zu", method->inputs_size(), method->outputs_size()); const size_t num_inputs = method_meta->num_inputs(); - std::vector> input_data(num_inputs); + std::vector> input_data(num_inputs); std::vector> input_sizes(num_inputs); std::vector> input_dim_order(num_inputs); std::vector> input_strides(num_inputs); @@ -326,8 +324,18 @@ int main(int argc, char** argv) { stride *= static_cast(input_sizes[i][d]); } - const size_t numel = static_cast(tensor_info->nbytes() / sizeof(float)); - input_data[i].assign(numel, 1.0f); + size_t numel = 1; + for (const auto size : input_sizes[i]) { + numel *= static_cast(size); + } + + input_data[i].assign(tensor_info->nbytes(), 0); + if (tensor_info->scalar_type() == exec_aten::ScalarType::Float) { + constexpr float one = 1.0f; + for (size_t value = 0; value < numel; ++value) { + std::memcpy(input_data[i].data() + value * sizeof(float), &one, sizeof(float)); + } + } fprintf(stderr, " input[%zu] shape=[", i); for (ssize_t d = 0; d < ndim; ++d) { diff --git a/py/torch-tensorrt-executorch-runtime/setup.py b/py/torch-tensorrt-executorch-runtime/setup.py index c78ffdfb5ad..5a022c590d8 100644 --- a/py/torch-tensorrt-executorch-runtime/setup.py +++ b/py/torch-tensorrt-executorch-runtime/setup.py @@ -137,7 +137,13 @@ def build_extension(self, ext: Extension) -> None: source = built.parent / dependency if not source.is_file(): raise RuntimeError(f"Bazel did not produce {source}") - shutil.copy2(source, output.parent / dependency) + destination = output.parent / dependency + # build_ext invokes this method once for each extension. Both + # extensions share these libraries, and copy2 preserves Bazel's + # read-only mode, so the second invocation cannot overwrite the + # first copy unless it is removed first. + destination.unlink(missing_ok=True) + shutil.copy2(source, destination) TENSORRT_DISTRIBUTION = tensorrt_distribution() diff --git a/tests/cpp/executorch/BUILD b/tests/cpp/executorch/BUILD index 17d2820bf25..4f7a3194b62 100644 --- a/tests/cpp/executorch/BUILD +++ b/tests/cpp/executorch/BUILD @@ -6,6 +6,7 @@ test_suite( name = "executorch_backend_tests", tests = [ ":test_caller_stream", + ":test_edge_llm_blob_header", ":test_executorch_binding_names", ":test_executorch_blob_header", ":test_executorch_weight_streaming_budget", @@ -39,6 +40,15 @@ cc_test( ], ) +cc_test( + name = "test_edge_llm_blob_header", + srcs = ["test_edge_llm_blob_header.cpp"], + deps = [ + "//cpp:edge_llm_executorch_blob_header", + "@googletest//:gtest_main", + ], +) + cc_test( name = "test_executorch_weight_streaming_budget", srcs = ["test_executorch_weight_streaming_budget.cpp"], diff --git a/tests/cpp/executorch/test_edge_llm_blob_header.cpp b/tests/cpp/executorch/test_edge_llm_blob_header.cpp new file mode 100644 index 00000000000..07afd97cebc --- /dev/null +++ b/tests/cpp/executorch/test_edge_llm_blob_header.cpp @@ -0,0 +1,103 @@ +#include "torch_tensorrt/executorch/EdgeLLMBlobHeader.h" + +#include "gtest/gtest.h" + +#include +#include +#include +#include +#include + +namespace torch_tensorrt { +namespace executorch_backend { +namespace { + +constexpr char EDGE_LLM_MAGIC[4] = {'E', 'L', '0', '1'}; +constexpr uint32_t HEADER_SIZE = 32; +constexpr uint32_t BLOB_ALIGNMENT = 16; + +template +void write_field(std::vector& payload, std::size_t offset, T value) { + std::memcpy(payload.data() + offset, &value, sizeof(value)); +} + +std::size_t align_up(std::size_t value, std::size_t alignment) { + return ((value + alignment - 1) / alignment) * alignment; +} + +std::string valid_metadata() { + return R"({"abi_version":1,"component":"vision","outputs":[{"dtype":"float16","shape":[1,4,8]}],)" + R"("runner":"vit","runner_config":{"model_type":"vit"}})"; +} + +std::vector make_payload(const std::string& metadata, const std::string& nested = "TR01nested") { + const auto metadata_offset = static_cast(HEADER_SIZE); + const auto metadata_size = static_cast(metadata.size()); + const auto blob_offset = static_cast(align_up(metadata_offset + metadata_size, BLOB_ALIGNMENT)); + std::vector payload(blob_offset + nested.size(), 0); + std::memcpy(payload.data(), EDGE_LLM_MAGIC, sizeof(EDGE_LLM_MAGIC)); + write_field(payload, 4, metadata_offset); + write_field(payload, 8, metadata_size); + write_field(payload, 12, blob_offset); + write_field(payload, 16, static_cast(nested.size())); + std::memcpy(payload.data() + metadata_offset, metadata.data(), metadata.size()); + std::memcpy(payload.data() + blob_offset, nested.data(), nested.size()); + return payload; +} + +TEST(ExecuTorchEdgeLLMBlobHeader, ParsesVisionPayload) { + const auto payload = make_payload(valid_metadata()); + + EdgeLLMBlobHeader header; + ASSERT_TRUE(EdgeLLMBlobHeader::parse(payload.data(), payload.size(), header)); + EXPECT_EQ(header.abi_version, 1); + EXPECT_EQ(header.component, "vision"); + EXPECT_EQ(header.runner, "vit"); + EXPECT_EQ(header.runner_config_json, R"({"model_type":"vit"})"); + EXPECT_EQ(header.blob_offset % BLOB_ALIGNMENT, 0); + EXPECT_EQ(EdgeLLMBlobHeader::nested_blob_data(payload.data(), header), payload.data() + header.blob_offset); +} + +TEST(ExecuTorchEdgeLLMBlobHeader, RejectsInvalidMagic) { + auto payload = make_payload(valid_metadata()); + payload[0] = 'X'; + EdgeLLMBlobHeader header; + EXPECT_FALSE(EdgeLLMBlobHeader::parse(payload.data(), payload.size(), header)); +} + +TEST(ExecuTorchEdgeLLMBlobHeader, RejectsUnsupportedAbi) { + auto metadata = valid_metadata(); + metadata.replace(metadata.find("\"abi_version\":1"), 15, "\"abi_version\":2"); + const auto payload = make_payload(metadata); + EdgeLLMBlobHeader header; + EXPECT_FALSE(EdgeLLMBlobHeader::parse(payload.data(), payload.size(), header)); +} + +TEST(ExecuTorchEdgeLLMBlobHeader, RejectsWrongRunner) { + auto metadata = valid_metadata(); + metadata.replace(metadata.find("\"runner\":\"vit\""), 14, "\"runner\":\"qwen\""); + const auto payload = make_payload(metadata); + EdgeLLMBlobHeader header; + EXPECT_FALSE(EdgeLLMBlobHeader::parse(payload.data(), payload.size(), header)); +} + +TEST(ExecuTorchEdgeLLMBlobHeader, RejectsEmptyOutputs) { + auto metadata = valid_metadata(); + const auto outputs = metadata.find("\"outputs\":["); + const auto outputs_end = metadata.find(']', outputs); + metadata.replace(outputs + 10, outputs_end - outputs - 9, "[]"); + const auto payload = make_payload(metadata); + EdgeLLMBlobHeader header; + EXPECT_FALSE(EdgeLLMBlobHeader::parse(payload.data(), payload.size(), header)); +} + +TEST(ExecuTorchEdgeLLMBlobHeader, RejectsNestedBlobPastPayload) { + auto payload = make_payload(valid_metadata()); + write_field(payload, 16, static_cast(payload.size())); + EdgeLLMBlobHeader header; + EXPECT_FALSE(EdgeLLMBlobHeader::parse(payload.data(), payload.size(), header)); +} + +} // namespace +} // namespace executorch_backend +} // namespace torch_tensorrt diff --git a/tests/py/dynamo/executorch/test_edge_llm_partitioner.py b/tests/py/dynamo/executorch/test_edge_llm_partitioner.py new file mode 100644 index 00000000000..66c38b1bb3b --- /dev/null +++ b/tests/py/dynamo/executorch/test_edge_llm_partitioner.py @@ -0,0 +1,205 @@ +from __future__ import annotations + +import sys +import types +from pathlib import Path + +import pytest +import torch +import torch.nn as nn + +pytest.importorskip("executorch.exir") + +# The regular exporters package eagerly registers optional model families. These +# focused lowering tests need only the package-local ops and must not require +# LeRobot or other model stacks. +if "exporters" not in sys.modules: + exporters_package = types.ModuleType("exporters") + exporters_package.__path__ = [ + str(Path(__file__).resolve().parents[4] / "tools/hf/exporters") + ] + sys.modules["exporters"] = exporters_package + +from exporters import ops as edge_ops +from exporters.executorch.artifact import build_vision_artifact +from exporters.executorch.backend import EdgeLLMBackend +from exporters.executorch.partitioner import EdgeLLMPartitioner +from exporters.executorch.serialization import ( + EdgeComponentMetadata, + EdgeOutputSpec, + deserialize_edge_component, + serialize_edge_component, +) +from exporters.executorch.vision import save_vision_pte +from torch_tensorrt.executorch.operator_support import TensorRTOperatorSupport +from torch_tensorrt.executorch.serialization import deserialize_engine + + +def _vision_metadata() -> EdgeComponentMetadata: + return EdgeComponentMetadata( + component="vision", + runner="vit", + outputs=(EdgeOutputSpec(shape=(1, 4, 8), dtype="float16"),), + runner_config={"model_type": "vit"}, + ) + + +class _VisionModule(nn.Module): + def __init__(self, metadata_json: str, trt_blob: bytes) -> None: + super().__init__() + self.metadata_json = metadata_json + self.register_buffer( + "trt_blob", torch.tensor(list(trt_blob), dtype=torch.uint8) + ) + + def forward(self, pixel_values): + return edge_ops.call_vision_tower( + self.trt_blob, self.metadata_json, pixel_values + )[0] + + +@pytest.mark.unit +def test_edge_component_payload_round_trip(): + metadata = _vision_metadata() + trt_blob = b"TR01-test-engine" + + payload = serialize_edge_component(trt_blob, metadata) + restored_blob, restored_metadata = deserialize_edge_component(payload) + + assert restored_blob == trt_blob + assert restored_metadata == metadata + + +@pytest.mark.unit +def test_vision_tower_fake_uses_embedded_output_spec(): + metadata = _vision_metadata() + fake_mode = torch._subclasses.fake_tensor.FakeTensorMode() + with fake_mode: + outputs = torch.ops.edge_llm.vision_tower.default( + [torch.empty(1, 16, 16, 3, device="cuda")], + torch.empty(16, dtype=torch.uint8), + metadata.to_json(), + ) + + assert len(outputs) == 1 + assert tuple(outputs[0].shape) == (1, 4, 8) + assert outputs[0].dtype == torch.float16 + assert outputs[0].device.type == "cuda" + + +@pytest.mark.unit +def test_edge_llm_partitioner_tags_only_vision_operator(): + metadata = _vision_metadata() + module = _VisionModule(metadata.to_json(), b"TR01-test-engine").eval() + exported = torch.export.export(module, (torch.randn(1, 16, 16, 3, device="cuda"),)) + + result = EdgeLLMPartitioner().partition(exported) + assert len(result.partition_tags) == 1 + spec = next(iter(result.partition_tags.values())) + assert spec.backend_id == EdgeLLMBackend.__name__ + + vision_node = next( + node + for node in result.tagged_exported_program.graph_module.graph.nodes + if node.op == "call_function" + and hasattr(node.target, "_schema") + and node.target._schema.name == "edge_llm::vision_tower" + ) + assert vision_node.meta["delegation_tag"] in result.partition_tags + assert not TensorRTOperatorSupport().is_node_supported({}, vision_node) + + +@pytest.mark.unit +def test_edge_llm_backend_wraps_vision_component(): + metadata = _vision_metadata() + trt_blob = b"TR01-test-engine" + module = _VisionModule(metadata.to_json(), trt_blob).eval() + exported = torch.export.export(module, (torch.randn(1, 16, 16, 3, device="cuda"),)) + + result = EdgeLLMBackend.preprocess(exported, []) + restored_blob, restored_metadata = deserialize_edge_component( + result.processed_bytes + ) + + assert restored_blob == trt_blob + assert restored_metadata == metadata + + +@pytest.mark.unit +def test_vision_metadata_rejects_wrong_runner(): + metadata = EdgeComponentMetadata( + component="vision", + runner="qwen_vit", + outputs=(EdgeOutputSpec(shape=(1, 4, 8), dtype="float16"),), + ) + with pytest.raises(ValueError, match="runner='vit'"): + edge_ops._vision_metadata(metadata.to_json()) + + +@pytest.mark.unit +def test_build_vision_artifact_wraps_saved_engine(tmp_path): + engine_dir = tmp_path / "vision" + engine_dir.mkdir() + (engine_dir / "visual.engine").write_bytes(b"serialized-vision-engine") + (engine_dir / "config.json").write_text("""{ + "model_type": "vit", + "component": "vision", + "engine_file": "visual.engine", + "input_names": ["pixel_values"], + "output_names": ["visual_embeds"], + "outputs": [{"shape": [4, 8], "dtype": "torch.float16"}], + "input_layout": "hwc", + "input_dtype": "float16" +} +""") + + artifact = build_vision_artifact(engine_dir, device_id=2) + nested_blob, edge_metadata = deserialize_edge_component( + serialize_edge_component( + artifact.trt_blob, + EdgeComponentMetadata.from_json(artifact.edge_metadata_json), + ) + ) + engine, trt_metadata = deserialize_engine(nested_blob) + + assert engine == b"serialized-vision-engine" + assert edge_metadata.runner_config["input_layout"] == "hwc" + assert trt_metadata.device_id == 2 + assert [binding.name for binding in trt_metadata.io_bindings] == [ + "pixel_values", + "visual_embeds", + ] + + +@pytest.mark.unit +def test_save_vision_pte_contains_edge_backend(tmp_path): + engine_dir = tmp_path / "vision" + engine_dir.mkdir() + (engine_dir / "visual.engine").write_bytes(b"serialized-vision-engine") + (engine_dir / "config.json").write_text("""{ + "model_type": "vit", + "component": "vision", + "engine_file": "visual.engine", + "input_names": ["pixel_values"], + "output_names": ["visual_embeds"], + "outputs": [{"shape": [4, 8], "dtype": "torch.float16"}], + "input_layout": "hwc", + "input_dtype": "float16" +} +""") + artifact = build_vision_artifact(engine_dir) + pte = tmp_path / "vision.pte" + + save_vision_pte( + artifact, + torch.randn(1, 16, 16, 3, device="cuda", dtype=torch.float16), + pte, + ) + + from executorch.exir._serialize._program import deserialize_pte_binary + + program = deserialize_pte_binary(pte.read_bytes()).program + delegate_ids = [ + delegate.id for plan in program.execution_plan for delegate in plan.delegates + ] + assert delegate_ids == ["EdgeLLMBackend"] diff --git a/tools/hf/export_pi05_vision_executorch.py b/tools/hf/export_pi05_vision_executorch.py new file mode 100644 index 00000000000..7573a8c2a2b --- /dev/null +++ b/tools/hf/export_pi05_vision_executorch.py @@ -0,0 +1,78 @@ +from __future__ import annotations + +import argparse +from pathlib import Path + +import torch +from exporters.compile import compile_component +from exporters.executorch import build_vision_artifact, save_vision_pte +from exporters.models.pi05.patches import PI05 +from exporters.models.pi05.vision import Pi05HwcVision +from exporters.plugin.attn_patches import apply_patches +from exporters.plugin.plugin_utils import load_plugins_for_trt +from exporters.spec import ComponentBundle +from exporters.utils import force_hf_attention +from lerobot.policies.pi05 import PI05Policy + +DEFAULT_CHECKPOINT = "lerobot/pi05_libero_base" + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Freshly compile PI0.5 vision and package EdgeLLMBackend .pte" + ) + parser.add_argument("--checkpoint", default=DEFAULT_CHECKPOINT) + parser.add_argument( + "--engine-dir", type=Path, default=Path("/tmp/pi05_executorch_fresh") + ) + parser.add_argument( + "--pte-path", type=Path, default=Path("/tmp/pi05_vision_edge.pte") + ) + parser.add_argument("--device-id", type=int, default=0) + args = parser.parse_args() + + device = torch.device("cuda", args.device_id) + dtype = torch.float16 + load_plugins_for_trt() + + policy = PI05Policy.from_pretrained(args.checkpoint).eval() + policy.model.to(device=device, dtype=dtype).eval() + paligemma = policy.model.paligemma_with_expert.paligemma.model + force_hf_attention(paligemma.vision_tower, "eager") + + pixel_values_hwc = torch.randn(4, 224, 224, 3, device=device, dtype=dtype) + bundle = ComponentBundle( + module=Pi05HwcVision(paligemma).eval(), + trace_args=(pixel_values_hwc,), + save_args=(pixel_values_hwc,), + input_names=["pixel_values"], + output_names=["visual_embeds"], + model_type="vit", + engine_file="visual.engine", + extra_config={ + "input_layout": "hwc", + "input_dtype": "float16", + }, + trt_settings={ + "disable_tf32": False, + "use_fp32_acc": False, + "use_explicit_typing": False, + "decompose_attention": True, + }, + ) + + with apply_patches(PI05): + engine_path, _, _ = compile_component( + bundle, + name="vision", + engine_dir=args.engine_dir, + ) + + artifact = build_vision_artifact(engine_path, device_id=args.device_id) + save_vision_pte(artifact, pixel_values_hwc, args.pte_path) + print(f"Saved fresh vision engine under {engine_path}") + print(f"Saved EdgeLLMBackend program to {args.pte_path}") + + +if __name__ == "__main__": + main() diff --git a/tools/hf/exporters/executorch/__init__.py b/tools/hf/exporters/executorch/__init__.py new file mode 100644 index 00000000000..0d6495770ba --- /dev/null +++ b/tools/hf/exporters/executorch/__init__.py @@ -0,0 +1,55 @@ +"""ExecuTorch lowering for named Edge-LLM component operators. + +Imports stay lazy so the core Edge exporter does not require ExecuTorch unless +the caller explicitly requests ExecuTorch lowering. +""" + +from __future__ import annotations + +from typing import Any + +__all__ = [ + "EDGE_LLM_ABI_VERSION", + "EDGE_LLM_MAGIC", + "EdgeComponentMetadata", + "EdgeExecuTorchArtifact", + "EdgeLLMBackend", + "EdgeLLMPartitioner", + "EdgeOutputSpec", + "build_vision_artifact", + "deserialize_edge_component", + "lower_vision_to_executorch", + "save_vision_pte", + "serialize_edge_component", +] + + +def __getattr__(name: str) -> Any: + if name == "EdgeLLMBackend": + from .backend import EdgeLLMBackend + + return EdgeLLMBackend + if name == "EdgeLLMPartitioner": + from .partitioner import EdgeLLMPartitioner + + return EdgeLLMPartitioner + if name in {"EdgeExecuTorchArtifact", "build_vision_artifact"}: + from . import artifact + + return getattr(artifact, name) + if name in {"lower_vision_to_executorch", "save_vision_pte"}: + from . import vision + + return getattr(vision, name) + if name in { + "EDGE_LLM_ABI_VERSION", + "EDGE_LLM_MAGIC", + "EdgeComponentMetadata", + "EdgeOutputSpec", + "deserialize_edge_component", + "serialize_edge_component", + }: + from . import serialization + + return getattr(serialization, name) + raise AttributeError(name) diff --git a/tools/hf/exporters/executorch/artifact.py b/tools/hf/exporters/executorch/artifact.py new file mode 100644 index 00000000000..cf639ed7f1d --- /dev/null +++ b/tools/hf/exporters/executorch/artifact.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from torch_tensorrt.executorch.serialization import ( + TensorRTBlobMetadata, + TensorRTIOBinding, + serialize_engine, +) + +from .serialization import EdgeComponentMetadata, EdgeOutputSpec + + +@dataclass(frozen=True) +class EdgeExecuTorchArtifact: + trt_blob: bytes + edge_metadata_json: str + + +def _dtype_name(value: Any) -> str: + return str(value).removeprefix("torch.") + + +def build_vision_artifact( + engine_dir: str | Path, + *, + device_id: int = 0, +) -> EdgeExecuTorchArtifact: + """Build the nested TR01 + Edge metadata for one VitRunner component.""" + component_dir = Path(engine_dir) + config_path = component_dir / "config.json" + try: + config = json.loads(config_path.read_text()) + if config["component"] != "vision" or config["model_type"] != "vit": + raise ValueError( + "Edge vision artifact requires component='vision' and model_type='vit'" + ) + input_names = list(config["input_names"]) + input_config = dict( + config.get("inputs", {}).get(input_names[0], {}) if input_names else {} + ) + input_shape = list(input_config.get("shape", [])) + input_layout = config.get("input_layout") + if input_layout is None and len(input_shape) == 4 and input_shape[-1] == 3: + input_layout = "hwc" + if input_layout != "hwc": + raise ValueError("Edge VitRunner artifact requires input_layout='hwc'") + output_names = list(config["output_names"]) + outputs = list(config["outputs"]) + engine_file = str(config["engine_file"]) + except (FileNotFoundError, json.JSONDecodeError, KeyError, TypeError) as exc: + raise ValueError(f"Invalid Edge vision config at {config_path}") from exc + + if len(input_names) != 1 or len(output_names) != 1 or len(outputs) != 1: + raise ValueError( + "The first Edge vision delegate requires exactly one input and one output" + ) + + engine_path = component_dir / engine_file + try: + engine_bytes = engine_path.read_bytes() + except FileNotFoundError as exc: + raise ValueError(f"Vision engine was not found at {engine_path}") from exc + + output_spec = EdgeOutputSpec.from_dict( + { + "shape": outputs[0]["shape"], + "dtype": _dtype_name(outputs[0]["dtype"]), + } + ) + trt_metadata = TensorRTBlobMetadata( + io_bindings=[ + TensorRTIOBinding(name=input_names[0], is_input=True), + TensorRTIOBinding( + name=output_names[0], + dtype=output_spec.dtype, + shape=list(output_spec.shape), + is_input=False, + ), + ], + device_id=device_id, + ) + edge_metadata = EdgeComponentMetadata( + component="vision", + runner="vit", + outputs=(output_spec,), + runner_config={ + "model_type": "vit", + "input_layout": "hwc", + "input_dtype": _dtype_name( + config.get("input_dtype", input_config.get("dtype", "float16")) + ), + }, + ) + return EdgeExecuTorchArtifact( + trt_blob=serialize_engine(engine_bytes, trt_metadata), + edge_metadata_json=edge_metadata.to_json(), + ) diff --git a/tools/hf/exporters/executorch/backend.py b/tools/hf/exporters/executorch/backend.py new file mode 100644 index 00000000000..0962ca800c4 --- /dev/null +++ b/tools/hf/exporters/executorch/backend.py @@ -0,0 +1,103 @@ +"""ExecuTorch backend preprocessing for Edge-LLM component delegates.""" + +from __future__ import annotations + +from typing import Any, final + +import torch +import torch.fx +from executorch.exir.backend.backend_details import ( + BackendDetails, + CompileSpec, + PreprocessResult, +) +from torch.export import ExportedProgram + +from .serialization import EdgeComponentMetadata, serialize_edge_component + +_VISION_SCHEMA = "edge_llm::vision_tower" + + +def _schema_name(target: Any) -> str: + if hasattr(target, "_schema"): + return str(target._schema.name) + return "" + + +def _vision_nodes(program: ExportedProgram) -> list[torch.fx.Node]: + return [ + node + for node in program.graph_module.graph.nodes + if node.op == "call_function" and _schema_name(node.target) == _VISION_SCHEMA + ] + + +def _resolve_tensor_constant( + program: ExportedProgram, node: torch.fx.Node +) -> torch.Tensor: + if node.op == "get_attr": + value = getattr(program.graph_module, node.target, None) + elif node.op == "placeholder": + target = node.target + for spec in program.graph_signature.input_specs: + arg = getattr(spec, "arg", None) + if arg is not None and getattr(arg, "name", None) == node.name: + target = spec.target or target + break + value = (program.state_dict or {}).get(target) + if value is None: + value = (program.constants or {}).get(target) + else: + raise ValueError( + f"Edge-LLM payload must be a constant tensor, got node op {node.op!r}" + ) + if not isinstance(value, torch.Tensor): + raise ValueError(f"Edge-LLM payload {node.name!r} did not resolve to a tensor") + return value + + +def _tensor_bytes(value: torch.Tensor) -> bytes: + data = value.detach().cpu().contiguous().view(torch.uint8) + return bytes(memoryview(data.numpy())) + + +@final +class EdgeLLMBackend(BackendDetails): # type: ignore[misc] + """Packs one named Edge component for the native EdgeLLMBackend.""" + + @staticmethod + def preprocess( + edge_program: ExportedProgram, + compile_specs: list[CompileSpec], + ) -> PreprocessResult: + del compile_specs + nodes = _vision_nodes(edge_program) + if len(nodes) != 1: + raise RuntimeError( + "EdgeLLMBackend expects exactly one vision_tower node per " + f"partition, found {len(nodes)}" + ) + + node = nodes[0] + if len(node.args) != 3: + raise RuntimeError( + "edge_llm::vision_tower must receive tensors, trt_blob, and " + f"metadata_json; found {len(node.args)} arguments" + ) + trt_blob_node = node.args[1] + metadata_json = node.args[2] + if not isinstance(trt_blob_node, torch.fx.Node): + raise ValueError("vision_tower trt_blob argument is not a graph value") + if not isinstance(metadata_json, str): + raise ValueError("vision_tower metadata_json argument must be a string") + + metadata = EdgeComponentMetadata.from_json(metadata_json) + if metadata.component != "vision" or metadata.runner != "vit": + raise ValueError( + "vision_tower payload must select component='vision' and " + f"runner='vit', got {metadata.component!r}/{metadata.runner!r}" + ) + + trt_blob = _resolve_tensor_constant(edge_program, trt_blob_node) + payload = serialize_edge_component(_tensor_bytes(trt_blob), metadata) + return PreprocessResult(processed_bytes=payload) diff --git a/tools/hf/exporters/executorch/operator_support.py b/tools/hf/exporters/executorch/operator_support.py new file mode 100644 index 00000000000..ddf2ff87ce2 --- /dev/null +++ b/tools/hf/exporters/executorch/operator_support.py @@ -0,0 +1,22 @@ +"""Operator support for the Edge-LLM ExecuTorch partitioner.""" + +from __future__ import annotations + +from typing import Dict + +import torch +from torch.fx.passes.operator_support import OperatorSupportBase + + +class EdgeLLMOperatorSupport(OperatorSupportBase): # type: ignore[misc] + """Recognizes named Edge component operators, starting with vision.""" + + _SUPPORTED_OPS = frozenset({"edge_llm::vision_tower"}) + + def is_node_supported( + self, submodules: Dict[str, torch.nn.Module], node: torch.fx.Node + ) -> bool: + del submodules + if node.op != "call_function" or not hasattr(node.target, "_schema"): + return False + return node.target._schema.name in self._SUPPORTED_OPS diff --git a/tools/hf/exporters/executorch/partitioner.py b/tools/hf/exporters/executorch/partitioner.py new file mode 100644 index 00000000000..146779e1556 --- /dev/null +++ b/tools/hf/exporters/executorch/partitioner.py @@ -0,0 +1,73 @@ +"""Partition named Edge-LLM component operators for ExecuTorch.""" + +from __future__ import annotations + +from typing import Callable, Dict, List, Optional, Tuple + +import torch +from executorch.exir.backend.compile_spec_schema import CompileSpec +from executorch.exir.backend.partitioner import ( + DelegationSpec, + Partitioner, + PartitionResult, +) +from executorch.exir.backend.utils import tag_constant_data +from torch.export import ExportedProgram +from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner + +from .backend import EdgeLLMBackend +from .operator_support import EdgeLLMOperatorSupport + +try: + from executorch.exir.passes.propagate_device_pass import ( + TARGET_DEVICE_COMPILE_SPEC_KEY as _TARGET_DEVICE_COMPILE_SPEC_KEY, + ) +except ImportError: + _TARGET_DEVICE_COMPILE_SPEC_KEY = "target_device" + + +class EdgeLLMPartitioner(Partitioner): # type: ignore[misc] + """Creates one EdgeLLMBackend partition per named component operator.""" + + def __init__( + self, + compile_specs: Optional[List[CompileSpec]] = None, + ) -> None: + super().__init__() + self.compile_specs = list(compile_specs or []) + if not any( + spec.key == _TARGET_DEVICE_COMPILE_SPEC_KEY for spec in self.compile_specs + ): + self.compile_specs.append( + CompileSpec(_TARGET_DEVICE_COMPILE_SPEC_KEY, b"cuda:0") + ) + self.delegation_spec = DelegationSpec( + backend_id=EdgeLLMBackend.__name__, + compile_specs=self.compile_specs, + ) + + def partition(self, exported_program: ExportedProgram) -> PartitionResult: + partitions = CapabilityBasedPartitioner( + exported_program.graph_module, + EdgeLLMOperatorSupport(), + allows_single_node_partition=True, + ).propose_partitions() + + partition_tags: Dict[str, DelegationSpec] = {} + for partition in partitions: + tag = f"edge_llm_{partition.id}" + for node in partition.nodes: + node.meta["delegation_tag"] = tag + partition_tags[tag] = self.delegation_spec + + tag_constant_data(exported_program) + return PartitionResult( + tagged_exported_program=exported_program, + partition_tags=partition_tags, + ) + + def ops_to_not_decompose( + self, ep: ExportedProgram + ) -> Tuple[List[torch._ops.OpOverload], Optional[Callable[[torch.fx.Node], bool]]]: + del ep + return ([], None) diff --git a/tools/hf/exporters/executorch/serialization.py b/tools/hf/exporters/executorch/serialization.py new file mode 100644 index 00000000000..0f2e5bd32a8 --- /dev/null +++ b/tools/hf/exporters/executorch/serialization.py @@ -0,0 +1,138 @@ +"""Versioned Edge-LLM component payloads for ExecuTorch delegates.""" + +from __future__ import annotations + +import json +import struct +from dataclasses import dataclass, field +from typing import Any + +EDGE_LLM_MAGIC = b"EL01" +EDGE_LLM_ABI_VERSION = 1 +HEADER_FORMAT = "<4sIIIQ8s" +HEADER_SIZE = struct.calcsize(HEADER_FORMAT) + + +def _align_to_16(offset: int) -> int: + return (offset + 15) & ~15 + + +@dataclass(frozen=True) +class EdgeOutputSpec: + shape: tuple[int, ...] + dtype: str + + @classmethod + def from_dict(cls, value: dict[str, Any]) -> "EdgeOutputSpec": + try: + shape = tuple(int(dim) for dim in value["shape"]) + dtype = str(value["dtype"]) + except (KeyError, TypeError, ValueError) as exc: + raise ValueError(f"Invalid Edge output specification: {value!r}") from exc + if any(dim < 0 for dim in shape): + raise ValueError( + f"Edge output shape must be static and non-negative: {shape}" + ) + if not dtype: + raise ValueError("Edge output dtype must not be empty") + return cls(shape=shape, dtype=dtype) + + def to_dict(self) -> dict[str, Any]: + return {"shape": list(self.shape), "dtype": self.dtype} + + +@dataclass(frozen=True) +class EdgeComponentMetadata: + component: str + runner: str + outputs: tuple[EdgeOutputSpec, ...] + runner_config: dict[str, Any] = field(default_factory=dict) + abi_version: int = EDGE_LLM_ABI_VERSION + + def validate(self) -> None: + if self.abi_version != EDGE_LLM_ABI_VERSION: + raise ValueError( + f"Unsupported Edge-LLM ABI version {self.abi_version}; " + f"expected {EDGE_LLM_ABI_VERSION}" + ) + if self.component not in {"vision", "language", "action"}: + raise ValueError(f"Unsupported Edge-LLM component {self.component!r}") + if not self.runner: + raise ValueError("Edge-LLM runner must not be empty") + if not self.outputs: + raise ValueError("Edge-LLM component must declare at least one output") + + def to_json(self) -> str: + self.validate() + return json.dumps( + { + "abi_version": self.abi_version, + "component": self.component, + "runner": self.runner, + "outputs": [output.to_dict() for output in self.outputs], + "runner_config": self.runner_config, + }, + separators=(",", ":"), + sort_keys=True, + ) + + @classmethod + def from_json(cls, value: str | bytes) -> "EdgeComponentMetadata": + try: + parsed = json.loads(value) + metadata = cls( + abi_version=int(parsed["abi_version"]), + component=str(parsed["component"]), + runner=str(parsed["runner"]), + outputs=tuple( + EdgeOutputSpec.from_dict(output) for output in parsed["outputs"] + ), + runner_config=dict(parsed.get("runner_config", {})), + ) + except (json.JSONDecodeError, KeyError, TypeError, ValueError) as exc: + raise ValueError("Invalid Edge-LLM component metadata") from exc + metadata.validate() + return metadata + + +def serialize_edge_component(trt_blob: bytes, metadata: EdgeComponentMetadata) -> bytes: + """Wrap a TR01/TR02 TensorRT blob in an EL01 component envelope.""" + metadata_json = metadata.to_json().encode("utf-8") + metadata_offset = HEADER_SIZE + payload_offset = _align_to_16(metadata_offset + len(metadata_json)) + header = struct.pack( + HEADER_FORMAT, + EDGE_LLM_MAGIC, + metadata_offset, + len(metadata_json), + payload_offset, + len(trt_blob), + b"\x00" * 8, + ) + padding = b"\x00" * (payload_offset - metadata_offset - len(metadata_json)) + return header + metadata_json + padding + trt_blob + + +def deserialize_edge_component( + payload: bytes, +) -> tuple[bytes, EdgeComponentMetadata]: + if len(payload) < HEADER_SIZE: + raise ValueError(f"Edge-LLM payload is too small: {len(payload)} bytes") + magic, metadata_offset, metadata_size, blob_offset, blob_size, _ = struct.unpack( + HEADER_FORMAT, payload[:HEADER_SIZE] + ) + if magic != EDGE_LLM_MAGIC: + raise ValueError(f"Invalid Edge-LLM payload magic: {magic!r}") + if metadata_offset < HEADER_SIZE: + raise ValueError("Edge-LLM metadata starts inside the payload header") + if blob_offset % 16 != 0: + raise ValueError("Nested TensorRT blob is not 16-byte aligned") + if metadata_offset + metadata_size > blob_offset: + raise ValueError("Edge-LLM metadata overlaps the nested TensorRT blob") + if blob_offset + blob_size > len(payload): + raise ValueError("Nested TensorRT blob extends past the Edge-LLM payload") + + metadata = EdgeComponentMetadata.from_json( + payload[metadata_offset : metadata_offset + metadata_size] + ) + return payload[blob_offset : blob_offset + blob_size], metadata diff --git a/tools/hf/exporters/executorch/vision.py b/tools/hf/exporters/executorch/vision.py new file mode 100644 index 00000000000..278eef3a20a --- /dev/null +++ b/tools/hf/exporters/executorch/vision.py @@ -0,0 +1,65 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import torch +import torch.nn as nn +from torch_tensorrt.executorch import export as export_executorch + +from ..ops import call_vision_tower +from .artifact import EdgeExecuTorchArtifact +from .partitioner import EdgeLLMPartitioner + + +class EdgeVisionModule(nn.Module): + """Export-only module containing one embedded Edge vision component.""" + + def __init__(self, artifact: EdgeExecuTorchArtifact) -> None: + super().__init__() + self.metadata_json = artifact.edge_metadata_json + self.register_buffer( + "trt_blob", + torch.frombuffer(bytearray(artifact.trt_blob), dtype=torch.uint8), + persistent=True, + ) + + def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: + return call_vision_tower(self.trt_blob, self.metadata_json, pixel_values)[0] + + +def lower_vision_to_executorch( + artifact: EdgeExecuTorchArtifact, + example_input_hwc: torch.Tensor, +) -> Any: + if example_input_hwc.ndim != 4: + raise ValueError( + "Edge VitRunner example input must have shape [B,H,W,C], got " + f"{tuple(example_input_hwc.shape)}" + ) + if example_input_hwc.dtype != torch.float16: + raise ValueError( + f"Edge VitRunner example input must be float16, got {example_input_hwc.dtype}" + ) + + exported = torch.export.export( + EdgeVisionModule(artifact).eval(), + (example_input_hwc,), + strict=False, + ) + return export_executorch( + exported, + partitioners=[EdgeLLMPartitioner()], + ) + + +def save_vision_pte( + artifact: EdgeExecuTorchArtifact, + example_input_hwc: torch.Tensor, + output_path: str | Path, +) -> None: + edge_program = lower_vision_to_executorch(artifact, example_input_hwc) + path = Path(output_path) + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("wb") as output: + edge_program.to_executorch().write_to_file(output) diff --git a/tools/hf/exporters/models/pi05/spec.py b/tools/hf/exporters/models/pi05/spec.py index c5f0bb74b69..4e84e8d6d09 100644 --- a/tools/hf/exporters/models/pi05/spec.py +++ b/tools/hf/exporters/models/pi05/spec.py @@ -25,6 +25,7 @@ pi05_compact_index, ) from .patches import PI05 +from .vision import Pi05HwcVision, nchw_to_hwc @register_edge_spec("pi05") @@ -281,17 +282,22 @@ def prepare( paligemma = core.paligemma_with_expert.paligemma.model language = paligemma.language_model px = sample["pixel_values"] + px_hwc = nchw_to_hwc(px) device = px.device dtype = px.dtype vision = ComponentBundle( - module=paligemma.eval(), - trace_args=(px,), - save_args=(px,), + module=Pi05HwcVision(paligemma).eval(), + trace_args=(px_hwc,), + save_args=(px_hwc,), input_names=["pixel_values"], output_names=["visual_embeds"], model_type="vit", engine_file="visual.engine", + extra_config={ + "input_layout": "hwc", + "input_dtype": str(px_hwc.dtype).removeprefix("torch."), + }, trt_settings={ "disable_tf32": False, "use_fp32_acc": False, @@ -451,7 +457,9 @@ def prepare( return {"vision": vision, "language": language_bundle, "action": action} def run(self, engines: Mapping[str, str], sample: Mapping[str, Any]) -> Any: - vis = call_engine(engines["vision"], "vision", sample["pixel_values"])[0] + vis = call_engine( + engines["vision"], "vision", nchw_to_hwc(sample["pixel_values"]) + )[0] prefix = fuse_prefix(vis, sample["lang_embeds"], sample["compact_index"]) lm = call_engine( engines["language"], diff --git a/tools/hf/exporters/models/pi05/vision.py b/tools/hf/exporters/models/pi05/vision.py new file mode 100644 index 00000000000..c990b3f9f15 --- /dev/null +++ b/tools/hf/exporters/models/pi05/vision.py @@ -0,0 +1,36 @@ +from __future__ import annotations + +import torch +import torch.nn as nn + + +def nchw_to_hwc(pixel_values: torch.Tensor) -> torch.Tensor: + if pixel_values.ndim != 4: + raise ValueError( + f"PI0.5 vision input must be rank 4, got {tuple(pixel_values.shape)}" + ) + return pixel_values.permute(0, 2, 3, 1).contiguous() + + +class Pi05HwcVision(nn.Module): + """Adapt VitRunner HWC input to the patched PI0.5 NCHW vision module.""" + + def __init__(self, paligemma: nn.Module) -> None: + super().__init__() + self.paligemma = paligemma + + def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: + if pixel_values.ndim != 4: + raise ValueError( + f"PI0.5 HWC vision input must be rank 4, got {tuple(pixel_values.shape)}" + ) + nchw = pixel_values.permute(0, 3, 1, 2).contiguous() + features = self.paligemma(nchw) + if features.ndim == 3: + features = features.reshape(-1, features.shape[-1]) + if features.ndim != 2: + raise ValueError( + "PI0.5 vision output must be rank 2 or 3 before flattening, " + f"got {tuple(features.shape)}" + ) + return features diff --git a/tools/hf/exporters/ops.py b/tools/hf/exporters/ops.py index a71f791827f..51e72703802 100644 --- a/tools/hf/exporters/ops.py +++ b/tools/hf/exporters/ops.py @@ -2,11 +2,17 @@ from __future__ import annotations +import hashlib +import json import sys +import threading +from pathlib import Path from typing import Any import torch +from .executorch.serialization import EdgeComponentMetadata + # One process-wide table so pytest dual-imports of this module still share # execute_engine state with record_engine. _REGISTRY: dict[str, Any] = sys.modules.setdefault( @@ -15,6 +21,10 @@ ) _ENGINE_META: dict[str, dict[str, Any]] = _REGISTRY["meta"] _COMPILED_MODULES: dict[str, torch.nn.Module] = _REGISTRY["modules"] +_ENGINE_LOAD_LOCK: threading.Lock = _REGISTRY.setdefault("lock", threading.Lock()) +_EMBEDDED_MODULES: dict[str, torch.nn.Module] = _REGISTRY.setdefault( + "embedded_modules", {} +) def record_engine( @@ -43,18 +53,60 @@ def _as_tuple(value: Any) -> tuple[torch.Tensor, ...]: return (value,) +def _load_serialized_engine(engine_path: str, component: str) -> torch.nn.Module: + from torch_tensorrt.dynamo.runtime import TorchTensorRTModule + + engine_dir = Path(engine_path) + config_path = engine_dir / "config.json" + try: + config = json.loads(config_path.read_text()) + engine_file = config["engine_file"] + input_names = config["input_names"] + output_names = config["output_names"] + except (FileNotFoundError, KeyError, json.JSONDecodeError) as exc: + raise RuntimeError( + f"Cannot load TensorRT engine metadata for {component!r} from " + f"{config_path}" + ) from exc + + serialized_path = engine_dir / engine_file + try: + serialized_engine = serialized_path.read_bytes() + except FileNotFoundError as exc: + raise RuntimeError( + f"Serialized TensorRT engine for {component!r} was not found at " + f"{serialized_path}" + ) from exc + + return TorchTensorRTModule( + serialized_engine=serialized_engine, + input_binding_names=list(input_names), + output_binding_names=list(output_names), + name=component, + ) + + +def _get_engine(engine_path: str, component: str) -> torch.nn.Module: + compiled = _COMPILED_MODULES.get(engine_path) + if compiled is not None: + return compiled + + # Engine deserialization is expensive and must only happen once per path. + with _ENGINE_LOAD_LOCK: + compiled = _COMPILED_MODULES.get(engine_path) + if compiled is None: + compiled = _load_serialized_engine(engine_path, component) + _COMPILED_MODULES[engine_path] = compiled + return compiled + + @torch.library.custom_op("edge_llm::execute_engine", mutates_args=()) # type: ignore[misc] def execute_engine( engine_path: str, component: str, tensors: list[torch.Tensor] ) -> list[torch.Tensor]: - compiled = _COMPILED_MODULES.get(engine_path) - if compiled is not None: - out = compiled(*tensors) - return list(_as_tuple(out)) - raise RuntimeError( - f"No in-process module for engine {engine_path!r} ({component}). " - "Compile first, or load a serialized engine at this path." - ) + compiled = _get_engine(engine_path, component) + out = compiled(*tensors) + return list(_as_tuple(out)) @execute_engine.register_fake # type: ignore[misc] @@ -81,6 +133,98 @@ def call_engine( return tuple(out) +def _tensor_bytes(value: torch.Tensor) -> bytes: + data = value.detach().cpu().contiguous().view(torch.uint8) + return bytes(memoryview(data.numpy())) + + +def _get_embedded_engine( + trt_blob: torch.Tensor, metadata: EdgeComponentMetadata +) -> torch.nn.Module: + from torch_tensorrt.dynamo.runtime import TorchTensorRTModule + from torch_tensorrt.executorch.serialization import deserialize_engine + + blob_bytes = _tensor_bytes(trt_blob) + cache_key = hashlib.sha256( + blob_bytes + metadata.to_json().encode("utf-8") + ).hexdigest() + compiled = _EMBEDDED_MODULES.get(cache_key) + if compiled is not None: + return compiled + + with _ENGINE_LOAD_LOCK: + compiled = _EMBEDDED_MODULES.get(cache_key) + if compiled is None: + engine_bytes, trt_metadata = deserialize_engine(blob_bytes) + input_names = [ + binding.name for binding in trt_metadata.io_bindings if binding.is_input + ] + output_names = [ + binding.name + for binding in trt_metadata.io_bindings + if not binding.is_input + ] + compiled = TorchTensorRTModule( + serialized_engine=engine_bytes, + input_binding_names=input_names, + output_binding_names=output_names, + name=f"edge_llm_{metadata.component}", + ) + _EMBEDDED_MODULES[cache_key] = compiled + return compiled + + +def _vision_metadata(metadata_json: str) -> EdgeComponentMetadata: + metadata = EdgeComponentMetadata.from_json(metadata_json) + if metadata.component != "vision" or metadata.runner != "vit": + raise ValueError( + "edge_llm::vision_tower requires component='vision' and runner='vit', " + f"got component={metadata.component!r}, runner={metadata.runner!r}" + ) + return metadata + + +@torch.library.custom_op("edge_llm::vision_tower", mutates_args=()) # type: ignore[misc] +def vision_tower( + tensors: list[torch.Tensor], + trt_blob: torch.Tensor, + metadata_json: str, +) -> list[torch.Tensor]: + metadata = _vision_metadata(metadata_json) + compiled = _get_embedded_engine(trt_blob, metadata) + return list(_as_tuple(compiled(*tensors))) + + +@vision_tower.register_fake # type: ignore[misc] +def _( + tensors: list[torch.Tensor], + trt_blob: torch.Tensor, + metadata_json: str, +) -> list[torch.Tensor]: + del trt_blob + metadata = _vision_metadata(metadata_json) + device = tensors[0].device if tensors else torch.device("cpu") + outputs = [] + for output in metadata.outputs: + dtype = getattr(torch, output.dtype, None) + if not isinstance(dtype, torch.dtype): + raise ValueError(f"Unsupported Edge output dtype {output.dtype!r}") + outputs.append(torch.empty(output.shape, dtype=dtype, device=device)) + return outputs + + +def call_vision_tower( + trt_blob: torch.Tensor, + metadata_json: str, + *tensors: torch.Tensor, +) -> tuple[torch.Tensor, ...]: + """Call the embedded PI0.5-compatible Edge-LLM vision component.""" + out = torch.ops.edge_llm.vision_tower.default( + list(tensors), trt_blob, metadata_json + ) + return tuple(out) + + @torch.library.custom_op("edge_llm::fuse_prefix", mutates_args=()) # type: ignore[misc] def fuse_prefix( vision_tokens: torch.Tensor, diff --git a/tools/hf/exporters/tests/test_edge_exporter.py b/tools/hf/exporters/tests/test_edge_exporter.py index 590ae602b29..af48fd196cb 100644 --- a/tools/hf/exporters/tests/test_edge_exporter.py +++ b/tools/hf/exporters/tests/test_edge_exporter.py @@ -7,7 +7,10 @@ import torch import torch.nn as nn import torch_tensorrt -from exporters import EdgeConfig, EdgeExporter, register_edge_spec +import torch_tensorrt.dynamo.runtime as trt_runtime +from exporters import EdgeConfig, EdgeExporter +from exporters import ops as exporter_ops +from exporters import register_edge_spec from exporters.ops import call_engine from exporters.spec import ComponentBundle, EdgeSpec, registered_specs from torch.export import ExportedProgram @@ -93,6 +96,66 @@ def test_edge_exporter_exported_program(tmp_path, monkeypatch): torch.testing.assert_close(out, expected) +@pytest.mark.unit +def test_execute_engine_prefers_in_process_module(tmp_path, monkeypatch): + engine_path = str(tmp_path / "language") + module = nn.Identity() + exporter_ops._COMPILED_MODULES[engine_path] = module + monkeypatch.setattr( + exporter_ops, + "_load_serialized_engine", + lambda *args: pytest.fail("serialized engine should not be loaded"), + ) + + try: + assert exporter_ops._get_engine(engine_path, "language") is module + finally: + exporter_ops._COMPILED_MODULES.pop(engine_path, None) + + +@pytest.mark.unit +def test_execute_engine_loads_and_caches_serialized_engine(tmp_path, monkeypatch): + engine_dir = tmp_path / "language" + engine_dir.mkdir() + (engine_dir / "language.engine").write_bytes(b"serialized-engine") + (engine_dir / "config.json").write_text("""{ + "engine_file": "language.engine", + "input_names": ["x"], + "output_names": ["y"] +} +""") + engine_path = str(engine_dir) + constructor_calls = [] + module = nn.Identity() + + def fake_runtime_module(**kwargs): + constructor_calls.append(kwargs) + return module + + monkeypatch.setattr( + trt_runtime, + "TorchTensorRTModule", + fake_runtime_module, + ) + + try: + first = exporter_ops._get_engine(engine_path, "language") + second = exporter_ops._get_engine(engine_path, "language") + finally: + exporter_ops._COMPILED_MODULES.pop(engine_path, None) + + assert first is module + assert second is module + assert constructor_calls == [ + { + "serialized_engine": b"serialized-engine", + "input_binding_names": ["x"], + "output_binding_names": ["y"], + "name": "language", + } + ] + + @pytest.mark.unit def test_attn_patch_attribute_restores(): from exporters.plugin.attn_patches import patch_attribute