OnnxRuntime EP Context 缓存功能设计

目录

背景

ONNX Runtime 的执行提供程序 (EPs) 使用户能够在后端 SDK（例如，QNN、OpenVINO、Vitis AI 等）的支持下，在各种硬件加速器上运行 ONNX 模型。
执行提供程序将 ONNX 模型转换为后端 SDK 所需的图格式，并将其编译成硬件所需的格式。
在 NPU 领域，这个转换和编译过程可能非常耗时，特别是对于 LLM 模型，有时需要数十分钟才能完成。这显著影响了会话创建期间的用户体验。
为了消除模型转换和编译的重复开销，大多数后端 SDK 提供了一项功能，可以将预编译模型导出到二进制文件中。

  • 后端 SDK 可以直接加载预编译模型并在目标设备上执行。
  • 这极大地缩短了会话创建时间并提高了整体效率。为了支持这项优化，ONNX Runtime 在 MS 域中引入了一个名为 EPContext 的贡献操作符。

EPContext 操作符 Schema

操作符域：com.microsoft
节点输入与输出：可变参数

属性表如下

属性 数据类型 描述
main_context int64 1（默认值）：此节点引用 EP 上下文内容，其中包含与此节点关联的图。
0：此节点不引用任何 EP 上下文内容。它期望从设置此字段为 1 的另一个节点中检索图。
一些 EP 支持包含多个图的单个上下文。
具有 main_context = 1 的 EPContext 节点指向主上下文。
此上下文包含多个图，这些图可以被设置 main_context = 0 的其他节点引用。
ep_cache_context string 如果 embed_mode = 1，则包含 EP 上下文的有效载荷；如果 embed_mode = 0，则包含上下文文件路径。
该路径相对于 ONNX 模型文件，可以是文件名或子文件夹/文件名。
embed_mode int64 1（默认值）：ep_cache_context 包含上下文内容的有效载荷。
0：ep_cache_context 包含上下文二进制文件的文件路径。
ep_sdk_version string 可选。用于生成此节点的 SDK 版本。
onnx_model_filename string 可选。原始 Onnx 模型文件名。
hardware_architecture string 可选。硬件架构。
partition_name string 可选。OnnxRuntime 分区图名称。
source string 可选。用于生成此节点的源标识符。
这应是由 EP 定义的唯一键，以支持 ONNX Runtime 运行使用不同 EP 的多个 EPContext 节点。
例如
QNN EP 只接受 source = QNN 或 QnnExecutionProvider 的节点。
OpenVINO EP 只接受 source = OpenVINOExecutionProvider 的节点。
notes string 可选。特定 EP 需要的附加信息。
max_size int64 可选。上下文中的最大大小，其用法取决于 EP。
默认值为 0。

EP Context node example

会话选项 描述
ep.context_enable 仅用于 EP context 模型生成。
1：使 ONNX Runtime 能够导出上下文缓存模型。
0（默认值）：禁用上下文模型导出。
ep.context_file_path 指定导出模型的文件路径。
默认值：对于上下文模型生成，为 original_file_name_ctx.onnx。
对于模型推理
如果用户从内存缓冲区加载模型且 EP context 二进制文件位于 ONNX 模型之外，则必须设置此选项。
ONNX Runtime EP 使用此路径来确定文件夹位置，并将其与 ep_cache_context（指向上下文二进制路径）结合，构建上下文二进制文件的绝对路径。
ep.context_embed_mode 仅用于上下文模型生成。
1：将 EP context 内容直接导出到 ONNX 模型中，存储在 ep_cache_context 节点属性内。
0（默认值）：将 EP context 内容导出到一个单独文件，并将文件名存储在 ONNX 模型中。
文件路径在 ep_cache_context 节点属性中跟踪。
ep.context_node_name_prefix 仅用于上下文模型生成。
指定 EPContext 节点名称的前缀（也用作 partition_name 属性和内部图名称）。
当多个 EPContext 节点组合成单个模型时，确保节点间的唯一性，防止命名冲突。
EP 还可以将此前缀应用于转换后的 EP context 二进制文件中的 ep_graph 名称。
session.model_external_initializers_file_folder_path 这并非 EPContext 设计特有。通常，对于包含外部数据的模型，当从内存缓冲区加载模型时，会话会丢失模型的名称和路径，导致无法找到外部数据文件。使用此配置可指定外部数据文件的文件夹路径。
所有外部数据文件应放置在同一文件夹中。
ep.context_model_external_initializers_file_name 仅用于上下文模型生成。
当某些节点在 CPU EP 上分区且这些节点具有外部初始化器时，使用此配置。生成 EP context 模型时，新模型不应依赖于源 ONNX 模型使用的旧外部数据文件。
当使用外部初始化器文件导出 EP context 模型时，使用此设置。
如果指定，所有初始化器将放置在外部数据文件中。
否则，所有初始化器将嵌入到生成的 ONNX 文件内部。
默认情况下，此选项未设置，这意味着所有初始化器都将包含在 ONNX 文件中。

EP Context 缓存模型生成工作流程

用于生成 EP Context 缓存模型的 EP Interface GetEpContextNodes()

直接在执行提供程序 (EP) 代码中生成分区图具有挑战性，因为 EP 无法完全了解整个分区图。为解决此问题，ONNX Runtime 引入了一个新的执行提供程序接口：GetEpContextNodes()。

virtual const InlinedVector<const Node*> GetEpContextNodes() const {
  return InlinedVector<const Node*>();
}
  • 此 API 返回一个指向 EPContext 节点的指针数组。
  • 如果执行提供程序需要生成上下文缓存模型，则应实现此接口。否则，可以不实现。
  • 创建 EPContext 节点及其依赖项（例如，如果 embed_mode = 0 则为上下文二进制文件）是 EP 的责任。
  • ONNX Runtime 图分区器使用此接口检索 EPContext 节点并生成分区 ONNX 模型。 EP context 模型生成代码详情在此

EP Context 缓存模型生成指南

OnnxRuntime EPs 应遵循以下指南来创建 EP context 缓存模型并保持统一的用户界面

  • 所有权
    • 执行提供程序 (EP) 负责创建 EPContext 节点及其依赖项。
    • ONNX Runtime 框架负责使用 EP 提供的 EPContext 节点列表生成 EP context ONNX 模型。
  • 生命周期
    • EPContext 节点的生命周期至少从 EP 调用编译时开始，到 EP 被销毁时结束。
  • ep.context_enable
    • 如果 ep.context_enable = 1，ONNX Runtime 会创建 EP context 缓存模型。
    • 否则，如果 ep.context_enable = 0（默认值），ONNX Runtime 将遵循标准工作流程，不生成缓存模型。
  • ep.context_file_path
    • 如果未提供 ep.context_file_path，ONNX Runtime 将通过将原始输入模型文件名中的 .onnx 替换为 _ctx.onnx 来生成输出模型文件名。
    • 如果指定了 ep.context_file_path，ONNX Runtime 将使用提供的文件路径。当 ep.context_embed_mode = 0 时，EP 也应使用此路径确定导出编译后的 EP context 二进制文件的文件夹位置。
    • 注意：当从内存缓冲区加载模型时，ep.context_file_path 是必需的，因为在这种情况下 ONNX Runtime 无法检索原始模型文件路径。
  • ep.context_embed_mode
    • 1：将 EP context 内容直接嵌入到 ONNX 模型中。
    • 0（默认值）：将 EP context 内容导出到单独文件（EP context 二进制文件）。
      • 即使存在多个分区子图，也应只有一个 EP context 二进制文件。如果 EP 短期内无法实现这一点，请在 EP 网页上注明。在这种情况下，用户将需要通过遍历所有主要 EPContext 节点（embed_mode=1 的节点）并从节点属性 ep_cache_context 中提取文件路径来确定生产部署所需的文件。
      • EP context 二进制文件名为 [model_name]_[ep].bin。
      • EP 在 EPContext 节点属性 ep_cache_context 中记录上下文二进制文件名。
      • 上下文二进制文件必须与导出的 ONNX 模型文件位于同一目录下。
      • EPContext 节点中记录的文件路径是相对于 ONNX 模型文件的相对路径。
      • 注意：允许使用子文件夹。
  • ep.context_node_name_prefix
    • 如果用户希望为 EPContext 节点名称添加自定义前缀（也应用于 partition_name 属性和图名称），EP 应在生成 EPContext 节点时提供此功能。
    • 当将来自不同模型的多个 EPContext 节点合并到单个模型中时，这非常有用，因为存在跨模型的节点名称或图名称冲突风险。
    • EP 应支持单个模型中的多个 EP context，使用户能够合并和互连从不同模型生成的 EPContext 节点。
  • 带外部数据的源模型
    当源模型依赖于外部数据文件时，ONNX 使用相对路径定位该文件。因此，外部数据文件必须与源模型位于同一目录。然而，新生成的模型不应依赖于任何原始源文件。此方法由以下几个考量驱动
    • 所有新生成的文件应位于同一目录。
    • 无法保证输出文件会生成在与源文件相同的目录中。
    • EPContext 设计允许模型由多个 EP 分区，每个 EP 编译自己的 EPContext 节点。统一和标准化的过程有助于避免数据重复。
    • 一些 EP 可能需要将权重从源文件复制到其 context 二进制文件中，以满足特定的数据布局要求。
    • 对于回退到 ONNX Runtime CPU EP 的子图，所有权重数据默认情况下将直接嵌入到新生成的 [model_name]_ctx.onnx 模型中。如果设置了 ep.context_model_external_initializers_file_name，则所有权重数据将保存到指定的外部初始化器文件中。

使用场景代码示例

通过从模型路径创建会话来生成 EPContext 模型

    Ort::SessionOptions so;

    // Enable EPContext ONNX model dumping
    so.AddConfigEntry(kOrtSessionOptionEpContextEnable, "1");

    // Add the execution provider (using QNN as an example)
    so.AppendExecutionProvider("QNN", provider_options);

    // Create the session to dump the `_ctx.onnx` model
    Ort::Session session1(env, "./model1.onnx", so);

通过从内存缓冲区中的模型创建会话来生成 EPContext 模型
与 C API CreateSessionFromArray 类似，下面的示例从存储在内存数组中的模型创建 ONNX Runtime 会话，导致会话无法跟踪模型的名称和路径。要生成 EPContext 模型，必须使用 ep.context_file_path 指定文件路径。

    // Read model file into buffer array
    std::vector<char> buffer;
    ReadFileToBuffer("./model1.onnx", buffer);

    Ort::SessionOptions so;

    // Enable EPContext ONNX model dumping
    so.AddConfigEntry(kOrtSessionOptionEpContextEnable, "1");

    // Specify the generated EPContext model file path using option ep.context_file_path
    so.AddConfigEntry(kOrtSessionOptionEpContextFilePath, "./model_ctx.onnx");

    // Add the execution provider (using QNN as an example)
    so.AppendExecutionProvider("QNN", provider_options);


    // Create the session to dump the `_ctx.onnx` model
    Ort::Session session1(env, buffer.data(), buffer.size(), so);

通过从内存缓冲区中的模型创建会话来生成 EPContext 模型，并且模型具有外部权重
从内存数组创建会话，并且模型依赖于外部数据。会话需要 session.model_external_initializers_file_folder_path 来确定外部数据位置，并且与前面的示例一样，需要 ep.context_file_path 来设置生成的 EPContext 模型的文件路径。

    // Read model file into buffer array
    std::vector<char> buffer;
    ReadFileToBuffer("./model_folder/model1.onnx", buffer);

    Ort::SessionOptions so;

    // Enable EPContext ONNX model dumping
    so.AddConfigEntry(kOrtSessionOptionEpContextEnable, "1");

    // Specify the generated EPContext model file path using option ep.context_file_path
    so.AddConfigEntry(kOrtSessionOptionEpContextFilePath, "./model_folder/model_ctx.onnx");

    // Specify the external data folder path using option session.model_external_initializers_file_folder_path
    so.AddConfigEntry(kOrtSessionOptionsModelExternalInitializersFileFolderPath, "./external_data_folder/");

    // Add the execution provider (using QNN as an example)
    so.AppendExecutionProvider("QNN", provider_options);


    // Create the session to dump the `_ctx.onnx` model
    Ort::Session session1(env, buffer.data(), buffer.size(), so);

注意：如果存在依赖于外部数据的 CPU EP 上的子图回退，生成的 EPContext 模型不应依赖于基础模型使用的原始外部数据文件。默认情况下，EPContext 模型将所有外部数据直接嵌入到生成的 ONNX 文件中。如果需要将权重存储在外部文件中，请设置 ep.context_model_external_initializers_file_name。此选项强制将所有初始化器保存到指定的外部文件中。

EP Context 缓存模型的推理工作流程

支持加载包含 EPContext 节点的模型的 ONNX Runtime EP 应遵循以下工作流程和规则进行模型推理

  • 模型识别
    • EP 应首先确定模型是否包含 EPContext 节点。
      • 如果不存在 EPContext 节点，EP 将遵循其正常的推理工作流程。
      • 如果模型包含 EPContext 节点
        • EP 应检查所有 EPContext 节点的 source 节点属性，以验证其中是否有任何节点是为当前 EP 设计的（即 source 属性与 EP 期望的键匹配）。
        • EP 应仅对 source 属性与 EP 所需的键匹配的 EPContext 节点进行分区。
        • EP 从匹配的 EPContext 节点加载缓存的 context。
  • 处理外部 Context 二进制文件 (embed_mode = 0) 当 EPContext 缓存模型使用 embed_mode = 0 生成时，Context 二进制文件作为单独的文件与 ONNX 模型一起存储在同一文件夹中。
    • ONNX Runtime 从 EPContext 节点的 ep_cache_context 属性中检索 Context 二进制文件的相对路径。
    • 对于从文件路径加载的模型
      • EP 应确定输入模型文件的文件夹路径，并将其与相对路径结合，构建到 Context 二进制文件的完整路径。
    • 对于从内存缓冲区加载的模型
      • 由于 EP 无法推导出模型的文件夹路径，用户必须指定会话选项 ep.context_file_path。
      • EP 使用 ep.context_file_path 确定文件夹路径，并将其与相对路径结合，构建到 Context 二进制文件的完整路径。
  • 支持多个主 EPContext 节点 (main_context = 1)
    • EP 应无限制地支持多个主 EPContext 节点。
    • EP 必须能够加载在 EPContext 节点的 ep_cache_context 属性中指定的所有 EP context 二进制缓冲区/文件，对其进行反序列化，管理 ep_graphs，并选择合适的进行执行。

  • EP Context 二进制加载期间的错误处理

    EP 或其后端 SDK 应能够检测常见的失败场景（包括但不限于以下情况）。在这种情况下，EP 应返回带有 INVALID_GRAPH 状态码的状态

    • 检测驱动程序版本与 EP context 二进制文件所需版本之间的不匹配；如果它们不兼容，则返回错误。
    • 检测运行时 SDK 版本与用于生成 EP context 二进制文件的版本之间的不匹配；如果它们不兼容，则返回错误。
    • 如果加载 EP context 二进制文件因任何原因失败，则返回错误。

EP Context nodes with different EPs

使用场景代码示例

从预编译的 EPContext 模型创建推理会话
从模型文件路径创建会话。如果存在外部 EP context 二进制文件，会话可以从模型文件路径中找到二进制文件路径。

    Ort::SessionOptions so;

    // Add EP, take QNN for example
    so.AppendExecutionProvider("QNN", provider_options);

    // Create sessions to load from the _ctx.onnx model
    Ort::Session session1(env, "model1_ctx.onnx", so);

    session1.run(...);

从内存缓冲区中的预编译 EPContext 模型创建推理会话
从模型的内存缓冲区创建会话会导致会话无法跟踪模型的名称和路径。要解决此问题，必须设置：ep.context_file_path。

  • 会话使用此路径来识别文件夹位置。
  • 利用 EPContext 节点中的 EP context 二进制文件名，会话构建到最终 EP context 二进制文件的完整路径。
      // Read model file into buffer array
  std::vector<char> buffer;
  ReadFileToBuffer("./model_folder/model_ctx.onnx", buffer);

  Ort::SessionOptions so;

  // Specify the EPContext model file path using option ep.context_file_path
  so.AddConfigEntry(kOrtSessionOptionEpContextFilePath, "./model_path/model_ctx.onnx");

  // Add EP, take QNN for example
  so.AppendExecutionProvider("QNN", provider_options);

  // Create sessions to load from the buffer
  Ort::Session session1(env, buffer.data(), buffer.size(), so);

  session1.run(...);

带权重共享的 EPContext

Onnx 域中的权重共享

在 ONNX 中，权重共享指的是多个具有外部权重的 ONNX 模型指向同一个外部权重文件。这些模型使用相同的张量名称，使它们能够引用相同的张量数据。

Weight sharing across Onnx models

在 EP 域中使用 EPContext 实现权重共享

EP 权重共享通过预生成的 EP context 二进制文件/ blob 实现。为此，用户必须离线（提前）生成 context 二进制文件。

  • 一些 EP 需要特定的平台，例如 Linux x86_64 和/或 Windows x86_64。详情请参考特定 EP 页面。
  • EP context 二进制文件包含共享相同张量的多个图。

Weight sharing in EP context binary

EP 或后端 SDK 应能够按照上述方式转换和编译图。

  • EP 或 SDK 应能从先前编译的图生成的现有 EP context 中识别出相同的权重。
  • 当新图编译到 EP context 中时，如果识别出相同的权重，应重用现有权重。例如，在 [model_name]_[ep].bin 中，来自 ep_graph1 的 tensor1_1 和来自 ep_graph2 的 tensor2_1 是相同的，并且都指向相同的数据偏移量 tensor_data1。

带权重共享的 EPContext 模型生成工作流程

Weight sharing workflow

每个 ONNX Runtime 会话都与一个 ONNX 模型关联。共享权重的模型被分组到模型组中，而具有共同属性的 ONNX Runtime 会话则被组织到会话组中。ONNX Runtime 引入了两个会话选项：ep.share_ep_contexts 和 ep.stop_share_ep_contexts，以促进会话分组。

  • 会话组内的所有 ONNX Runtime 会话都应启用 ep.share_ep_contexts。
  • 最后一个 ONNX Runtime 会话使用 ep.stop_share_ep_contexts 指示它是组中的最后一个会话。注意：单个 ONNX 模型可能包含多个 EPContext 节点，具体取决于图分区结果。但是，为简单起见，此处每个模型仅显示一个 EPContext 节点。

带权重共享的 EPContext 模型生成实现指南

  • 共享工作空间创建
    第一个会话创建共享工作空间（例如，EP 单例）以与其他会话共享资源。
  • EP Context 二进制文件命名
    EP context 二进制文件名由第一个会话确定，并存储在共享工作空间（例如，EP 单例）中，供跨会话组使用。
    EP context 二进制文件名应为 [model1_name]_[ep].bin。
  • 图编译
    会话组中的所有会话将其图编译到共享资源中。
  • EPContext 模型生成
    会话组中的每个会话都会创建一个 EPContext ONNX 模型。EP 生成一个引用 EP context 二进制文件名的 EPContext 节点。然后，ONNX Runtime 框架导出 EPContext ONNX 模型。
  • 最终 EP Context 二进制文件生成
    会话组中的最后一个会话（启用 ep.stop_share_ep_contexts 的会话）使用存储在共享工作空间中的名称生成最终的 EP context 二进制文件。
  • 共享工作空间清理
    最后一个会话清理共享工作空间。空的共享工作空间表示下一个运行的会话是第一个会话。
  • 生成的文件数量
    对于 N 个共享权重的源模型，应总共生成 N+1 个文件。
    生成的文件是 model1_ctx.onnx, ..., modeln_ctx.onnx, [model1_name]_[ep].bin。

用户代码示例

    Ort::SessionOptions so;

    // Enable EPContext ONNX model dumping
    so.AddConfigEntry(kOrtSessionOptionEpContextEnable, "1");

    // Enable EP context sharing across sessions
    so.AddConfigEntry(kOrtSessionOptionShareEpContexts, "1");

    // Add the execution provider (using QNN as an example)
    so.AppendExecutionProvider("QNN", provider_options);

    // Create the first session to dump the model1_ctx.onnx file
    Ort::Session session1(env, "model1.onnx", so);

    // Mark the last session by enabling ep.stop_share_ep_contexts
    so.AddConfigEntry(kOrtSessionOptionStopShareEpContexts, "1");

    // Create the last session to dump the model2_ctx.onnx file and generate the [model1_name]_[ep].bin
    Ort::Session session2(env, "model2.onnx", so);

带权重共享的 EPContext 模型生成通用工具

OnnxRuntime 提供了 ep_weight_sharing_ctx_gen 工具来自动化权重共享工作流程。该工具处理整个过程。此工具专为权重共享场景设计，简化了 EPContext 模型生成过程。示例命令行：

./ep_weight_sharing_ctx_gen -e qnn -i "soc_model|60 htp_graph_finalization_optimization_mode|3" ./model1.onnx,./model2.onnx

它会创建两个 Onnx 模型 (model1_ctx.onnx, model2_ctx.onnx) 和一个 QNN 上下文二进制文件 ([model1_name]_[ep].bin)。

通过启用权重共享的 EPContext 模型进行推理会话

要使用导出的启用权重共享的 EPContext 模型，ONNX Runtime 推理会话必须激活资源共享。这通过设置会话选项实现：

    ep.share_ep_contexts = 1

通过启用权重共享的 EPContext 模型进行推理的实现指南

  • 创建第一个 OnnxRuntime 推理会话
    • 设置会话选项：ep.share_ep_contexts=1。
    • 加载 model1_ctx.onnx 模型。
    • 共享工作空间最初为空。
    • EP 加载 [model1_name]_[ep].bin 并反序列化二进制文件以检索所有图（例如，ep_graph1, ep_graph2）。
    • model1_ctx.onnx 中的 EPContext 节点指定使用 ep_graph1。
    • 会话使用 ep_graph1 进行推理。
    • 剩余的图 (ep_graph2) 被放入共享工作空间供未来的会话使用。
  • 创建第二个 ONNX Runtime 推理会话
    • 设置会话选项：ep.share_ep_contexts=1。
    • 加载 model2_ctx.onnx 模型。
    • model2_ctx.onnx 中的 EPContext 节点指定使用 ep_graph2。
    • 共享工作空间已包含 ep_graph2。
    • 由于所需的图已在共享工作空间中可用，EP 跳过加载 [model1_name]_[ep].bin。
    • 会话将 ep_graph2 从共享工作空间移至当前会话，使其在共享工作空间中不再可访问。
  • 会话清理最佳实践
    • 为避免并发执行期间出现问题，建议按相反顺序销毁会话（即，先销毁第二个会话，再销毁第一个会话）。
    • 这确保了适当的资源管理，并防止与共享资源的潜在冲突。

用户代码示例

    Ort::SessionOptions so;
    // enable ep.share_ep_contexts
    so.AddConfigEntry(kOrtSessionOptionShareEpContexts, "1");

    // Add EP, take QNN for example
    so.AppendExecutionProvider("QNN", provider_options);

    // Create sessions to load from the _ctx.onnx models with resource sharing enabled
    Ort::Session session1(env, "model1_ctx.onnx", so);	
    Ort::Session session2(env, "model2_ctx.onnx", so);

    session1.run(...);
    session2.run(...);