当前位置：首页 > news >正文

昇腾CANN triton-inference-server-ge-backend：Triton 推理服务在 NPU 上的部署实战

news 2026/5/23 5:09:55

Triton Inference Server 是 NVIDIA 开源的推理服务器，支持多模型、多框架、动态 batch、模型版本管理。triton-inference-server-ge-backend 是 CANN 社区的适配层——把 Triton 的后端从 TensorRT/CUDA 换成 GE（Graph Engine），让 Triton 跑在昇腾 NPU 上。

适配架构

Triton 的模型服务分三层，这个仓库只动最底层——backend：

┌─────────────────────────────────────────┐ │ Triton Server（框架无关） │ │ HTTP/gRPC API + 负载均衡 + 模型版本管理 │ ├─────────────────────────────────────────┤ │ Backend API（标准接口） │ │ TRITONBACKEND_ModelInitialize() │ │ TRITONBACKEND_ModelInstanceExecute() │ │ TRITONBACKEND_ModelFinalize() │ ├─────────────────────────────────────────┤ │ Backend 实现（这里做替换） │ │ ┌──────────┬──────────┬──────────────┐ │ │ │ TensorRT │ ONNX RT │ GE Backend ◀─┼── 这个仓库 │ │ (GPU) │ (GPU) │ (NPU) │ │ │ └──────────┴──────────┴──────────────┘ │ └─────────────────────────────────────────┘

GE Backend 的核心逻辑：接收 Triton 的推理请求 → 把输入 tensor 转为 GE 格式 → 调 GE Graph Executor 执行 → 输出 tensor 返回给 Triton。

// triton-inference-server-ge-backend/src/ge_backend.cc// Triton Backend API 的三个标准接口实现TRITONSERVER_Error*TRITONBACKEND_ModelInitialize(TRITONBACKEND_Model*model){// 第一步：读取模型配置（Triton config.pbtxt）constchar*model_path;TRITONBACKEND_ModelPath(model,&model_path);// 第二步：用 GE 加载 .om 模型文件// .om 是离线编译好的模型（ATC 编译器输出）ge::Session*session=newge::Session(ge_options);session->LoadModel(model_path+"/model.om");// 第三步：写模型 state 到 Triton（后面执行时用到）TRITONBACKEND_ModelSetState(model,session);returnnullptr;// 成功}TRITONSERVER_Error*TRITONBACKEND_ModelInstanceExecute(TRITONBACKEND_ModelInstance*instance,TRITONBACKEND_Request**requests,uint32_trequest_count){// 获取初始化时保存的 sessionge::Session*session;TRITONBACKEND_ModelState(TRITONBACKEND_ModelInstanceModel(instance),(void**)&session);// 遍历所有推理请求for(uint32_tr=0;r<request_count;r++){// 从 Triton 请求中提取输入 tensor// Triton 的 tensor → GE 的 TensorDescstd::vector<ge::Tensor>inputs,outputs;TritonToGeTensors(requests[r],inputs,outputs);// 执行推理：GE Graph Executorge::Status status=session->Run(inputs,outputs);// 输出 tensor 写回 TritonGeToTritonOutputs(outputs,requests[r]);}returnnullptr;}TRITONSERVER_Error*TRITONBACKEND_ModelFinalize(TRITONBACKEND_Model*model){ge::Session*session;TRITONBACKEND_ModelState(model,(void**)&session);session->UnloadModel();deletesession;returnnullptr;}

三个接口——初始化（加载模型）、执行（推理）、销毁——就是 GE Backend 的全部代码。适配工作不是写算子，是把两个现有系统（Triton 和 GE）的接口对接起来。

模型部署流程

完整流程分三步：PyTorch 模型 → ATC 编译 → Triton 配置 → 启动服务。

# 第一步：PyTorch → ONNX → .om（离线编译）python torch_to_om.py\--modelllama-7b\--input_shape"1,512"\--outputllama-7b.om# 第二步：Triton 模型仓库目录结构# model_repository/# llama-7b/# config.pbtxt ← Triton 配置# 1/# model.om ← GE Backend 加载这个文件# version# config.pbtxtname:"llama-7b"backend:"ge"max_batch_size:8input[{name:"input_ids", data_type: TYPE_INT64, dims:[-1]}]output[{name:"logits", data_type: TYPE_FP32, dims:[-1,32000]}]instance_group[{count:1, kind: KIND_CPU}# .om 文件本身在 NPU 上跑]# 第三步：启动 Tritontritonserver\--model-repository=/model_repository\--backend-directory=/opt/tritonserver/backends\--grpc-port=8001\--http-port=8000# 验证服务curlhttp://localhost:8000/v2/models/llama-7b# 返回模型信息和状态

动态 batch 与 GE 的配合

Triton 最核心的功能是动态 batch——把同一时间到达的多个推理请求合并成一个大 batch，摊薄 kernel launch 的开销。GE Backend 支持这个功能，但有一些限制：

// GE Backend 中的 batch 处理逻辑// Triton 传递的 batch_size 是动态的（1-8 可变）// 但 .om 模型是固定 shape 编译的（比如 [1, 512]）// GE Backend 需要做 padding：把 batch < max_batch 的请求 pad 到固定 batchTRITONSERVER_Error*ExecuteBatch(ge::Session*session,TRITONBACKEND_Request**requests,uint32_trequest_count// 实际 batch){// .om 编译时固定 batch=8constintfixed_batch=8;// 构造固定大小的输入 tensorge::Tensorinput_tensor({fixed_batch,max_seq_len},ge::DATA_TYPE_INT64);// 拷贝真实数据（前 request_count 个）for(inti=0;i<request_count;i++){CopyTritonInput(requests[i],input_tensor,offset=i);}// 剩余位置用 pad token 填充// 注意：pad token 不影响 attention mask// mask 仍然是 [request_count, max_seq_len]，不是 [fixed_batch, max_seq_len]for(inti=request_count;i<fixed_batch;i++){FillPadding(input_tensor,pad_token_id,offset=i);}// 推理ge::Tensor output_tensor;session->Run({input_tensor},{output_tensor});// 只返回有效结果（前 request_count 个）for(inti=0;i<request_count;i++){CopyGeOutput(output_tensor,requests[i],offset=i);}}

踩坑一：.om 编译时的 batch 大小定死了

ATC 编译的 .om 模型是固定 batch 的。如果编译时 batch=4，Triton 配置 max_batch_size=8 → 推理时传入 batch=8 的请求，GE 执行报 shape mismatch。

错误配置：

# 编译时 batch=4atc--model=llama-7b.onnx--output=llama-7b.om\--input_format=ND--input_shape="input_ids:4,512"# Triton 配置 max_batch_size=8 → 冲突！# GE 在 Run() 时检测到 input 的 batch 维度是 8# 但 .om 编译时固定为 4 → shape mismatch 异常

正确配置：Triton 的 max_batch_size 必须 ≤ .om 编译时的固定 batch。

# 先把 max_batch 设为 8 编译atc--model=llama-7b.onnx--output=llama-7b.om\--input_format=ND--input_shape="input_ids:8,512"# Triton 配置 max_batch_size ≤ 8# config.pbtxtmax_batch_size:8

踩坑二：GE Session 不是线程安全的

Triton 的模型实例可以多线程并发执行。但 GE Session::Run() 不是线程安全的——两个线程同时调 Run() 会触发 GE 内部的 graph reorder 冲突。

错误配置：

// config.pbtxt instance_group [ { count: 4, // 4 个实例 kind: KIND_GPU // 但 GE Session 不能多线程共享 } ] // 4 个实例并发调用 session->Run() // GE 内部触发 graph reorder → 两个线程跑的时候 graph 结构被改了 // 报错：GraphExecutor: graph st // GE 内部触发 graph reorder → 两个线程跑的时候 graph 结构被改了 // 报错：GraphExecutor: graph status is executing

正确配置：每个 instance 创建独立的 GE Session。

// 正确：每个 Triton instance 有自己的 sessionTRITONBACKEND_ModelInitialize(model){for(inti=0;i<instance_count;i++){ge::Session*session=newge::Session(options);session->LoadModel(model_path);instance_sessions.push_back(session);}TRITONBACKEND_ModelSetState(model,instance_sessions);}// 执行时每个线程取各自的 sessionge::Session*my_session=instance_sessions[instance_id];my_session->Run(inputs,outputs);// 各自独立，无冲突

踩坑三：Triton 的 KV Cache 跨请求不共享

Triton 的默认行为是每个请求独立执行——请求之间不共享 KV Cache。但大模型推理的 KV Cache 复用是延迟优化的关键——第 N 个 token 的推理可以复用前 N-1 个 token 的 KV Cache。

GE Backend 需要显式支持 prefix caching：

// KV Cache 复用实现（简化）classKVCacheManager{// 以 prompt hash 为 key，缓存整段 prompt 的 KVstd::unordered_map<std::string,ge::Tensor>cache_;public:ge::Tensor*Lookup(conststd::string&prompt_hash){autoit=cache_.find(prompt_hash);return(it!=cache_.end())?&(it->second):nullptr;}voidStore(conststd::string&prompt_hash,ge::Tensor kv_cache){// FIFO 淘汰：最多缓存 16 个 prompt 的 KVif(cache_.size()>=16){autooldest=cache_.begin();cache_.erase(oldest);}cache_[prompt_hash]=kv_cache;}};

系统 prompt 在多轮对话中完全不变——缓存一次，后续所有请求都复用。第一次请求 512 token 的延迟是 200ms，后续同样的 prompt 只有 50ms（纯 KV Cache lookup）。

triton-inference-server-ge-backend 看似简单（几百行 C++），但每一行都在桥接两个生态的语义差异：Triton 的推理模型和 GE 的静态图、Triton 的多后端架构和 NPU 的独占执行模式、Triton 的动态 batch 和 ATC 的固定 shape。适配层不需要写复杂算法——需要对两个系统的边界条件足够清楚。

查看全文

http://www.jsqmd.com/news/868833/