大模型推理层服务化架构

模型推理服务化架构

一个Spring Boot Java进程，内部通过JNI加载了ONNX Runtime的native库，该库将模型文件读入Java进程的堆外内存中，然后在这个进程内（用显卡计算）同步/异步执行推理。

graph TB
    subgraph 客户端
        Client["客户端<br>(Web/App/微服务)"]
    end

    subgraph 网关层
        APIGateway["API Gateway<br>(Kong/APISIX/Nginx)"]
        Monitor["监控系统<br>(Prometheus + Grafana)"]
    end

    subgraph "Java推理服务(Spring Boot)"
        Controller[REST Controller]
        PrePost[预处理/后处理<br>（token化、归一化、缩放等）]
        ORTSession[ONNX Runtime Session<br>（JNI代理）]
        NativeMemory[堆外内存<br>模型权重 + 中间缓冲区]
    end

    subgraph "本地硬件"
        NativeLib["ONNX Runtime Native Lib<br>(.dll / .so)"]
        GPU["GPU (CUDA)"]
        ModelFile["模型文件<br>(.onnx)"]
    end

    Client -->|HTTPS| APIGateway
    APIGateway -->|路由/限流| Controller
    Controller --> PrePost
    PrePost -->|预处理后的张量| ORTSession
    ORTSession -.->|JNI调用| NativeLib
    NativeLib -->|读取| ModelFile
    NativeLib -->|加载到| NativeMemory
    NativeLib -->|CUDA执行| GPU
    GPU -->|推理结果| NativeLib
    NativeLib -->|返回张量| ORTSession
    ORTSession --> PostProcess[后处理]
    PostProcess -->|最终响应| Controller
    Controller --> APIGateway
    APIGateway --> Client

    ORTSession -.->|指标暴露| Monitor
    Controller -.->|业务指标| Monitor
    Monitor -.->|告警/大盘| APIGateway

Client → API Gateway （new-api之类） → Spring Boot Service（Java进程：Model加载到其堆外内存，JNI调用本机安装的native ONNXRuntime用显卡去跑堆外内存的Model）

• ONNX Runtime的C++核心库通过JNI被调用，模型权重、中间激活值都分配在Native Memory，不经过JVM堆，避免GC压力。

• Java API底层就是JNI调用预编译的libonnxruntime.so（Linux）或.dll（Windows）或.dylib（macOS）。这些动态库直接与硬件交互。

• ONNX Runtime支持CUDA（NVIDIA GPU）、DirectML、OpenVINO等执行后端。需要：

  • 安装对应版本的onnxruntime-gpu（Maven坐标：com.microsoft.onnxruntime:onnxruntime_gpu）。
  • 确保系统有NVIDIA驱动、CUDA库。
  • 在创建session时通过OrtSession.SessionOptions 指定设备ID（比如options.addCUDA(0)）。
  • 模型推理时，数据会从CPU内存拷贝到GPU显存，计算在GPU上完成，结果再拷回CPU内存（这个过程ORT会自动管理）。

统一的模型文件

ONNX Runtime .onnx模型文件

ONNX Runtime 的 Java API 会通过 JNI 调用它本地的 C++ 库（比如 libonnxruntime.so 或 onnxruntime.dll），这个C++库将模型的计算图和权重数据加载到 Java 进程的堆外内存（native memory）中。

而.onnx 模型文件是统一的模型文件，转化路径如下：

graph TD
    subgraph PyTorch
        A["PyTorch模型<br>(.pth 或 .pt)"] --> B[torch.onnx.export]
    end
    subgraph TensorFlow
        C["TensorFlow模型<br>(.h5 或 .pb)"] --> D[tf2onnx]
    end
    B --> E["统一ONNX模型<br>(.onnx)"]
    D --> E
    E --> F[部署到ONNX Runtime]

另外除了PyTorch 和 TensorFlow，实际上还有一条很常见的路径：

Hugging Face 模型 → optimum 库 → .onnx

现在大量的 NLP/embedding 模型都在 HuggingFace 上，optimum 是专门做这个转换的，而且会做一些针对推理的图优化，比直接用 torch.onnx.export 出来的模型在推理性能上通常更好。

更大规模的推理层

现在的架构是 Spring Boot 内嵌 ORT，这在中小规模完全合理。但如果模型变大、请求量上来，业界还有另一种做法是把推理服务单独拆出去，用 Triton Inference Server（NVIDIA 出的，支持 ONNX、TensorRT、TorchScript 等多种格式），Java 服务通过 gRPC 调它。这样的好处是推理层可以独立扩缩容，GPU 资源利用率更高，而且 Triton 自带 dynamic batching 和 model ensemble。