本文基于昇腾CANN和昇腾NPU围绕 cann-recipes-infer 仓库的相关技术展开。AI Agent 的推理模式跟普通对话不一样模型输出的是一个动作序列——调工具、看结果、再调工具直到任务完成。这对推理系统的要求从一次预测变成了循环直到停止。CANN 在 Agent 推理场景下需要做两件事高吞吐的循环推理 工具调用的 Host-Device 交互。Agent 推理循环的结构# AI Agent 的典型推理循环——工具调用直到完成classAgentExecutor: 一个 Agent 推理单元模型 工具集合 停止条件 def__init__(self,model,tools):self.modelmodel# LLM 模型self.tools{search:ToolSearch(),calc:ToolCalculator(),code:ToolCodeRunner(),read:ToolFileReader(),}self.max_turns10self.history[]defrun(self,user_query): 一次用户请求 N 轮模型推理 M 次工具调用 messages[{role:user,content:user_query}]forturninrange(self.max_turns):# Step 1: 模型推理——生成回复responseself.model.chat(messages)# Step 2: 解析是否要调工具# Agent 模型的输出格式通常是# tool_call{name: search, args: {query: ...}}/tool_calltool_callsparse_tool_calls(response)ifnottool_calls:# 没有工具调用 → Agent 完成了returnresponse# Step 3: 执行工具调用——Host 侧fortcintool_calls:toolself.tools[tc[name]]resulttool.run(tc[args])# 把工具结果追加给模型messages.append({role:tool,tool_call_id:tc[id],content:result})# Step 4: 继续下一轮——带着工具结果再次推理returnAgent 达到最大轮数未完成Agent 推理的核心约束每次工具调用完KV Cache 要能续上——不能每次重新 Prefill。KV Cache 续接——避免重复 Prefill# Agent 推理的关键优化KV Cache 续接classAgentKVCacheManager: 管理 Agent 多轮推理的 KV Cache——不走重复 Prefill def__init__(self,num_layers,num_heads,head_dim,max_seq65536):# 预分配一个大的 KV Cache 空间self.kv_pooltorch.empty(num_layers,2,max_seq,num_heads,head_dim,dtypetorch.float16,devicenpu)self.current_len0defextend_with_tool_result(self,tool_result_tokens): 工具调用结果追加到已有 Cache 后面 流程 Prefill 阶段首次Query → 模型 Cache[0..q_len) 工具调用后Cache[0..q_len) 不动新 Cache[q_len..) 追加 start_posself.current_len end_posstart_poslen(tool_result_tokens)# 在 CANN 上做增量 Prefill# 不是从头算全部只算新增部分的 K、V# 模型的其他层 Cache 不动new_kvself.incremental_prefill(tool_result_tokens,start_pos# 告诉模型从哪个位置开始)# 写到预先分配的位置forlayerinrange(num_layers):self.kv_pool[layer,0,start_pos:end_pos]new_kv[layer,0]self.kv_pool[layer,1,start_pos:end_pos]new_kv[layer,1]self.current_lenend_posreturnstart_posdefincremental_prefill(self,tokens,start_pos): 增量 Prefill——只算新增 Token 的 K、V CANN Runtime 支持这种从指定位置开始的执行模式 不需要重新构建整图的 KV Cache # CANN 上通过 GE 的 Incremental 执行模式# 设置输入位置偏移模型自动从 start_pos 开始outputmodel.execute_with_offset(input_idstokens,position_offsetstart_pos,kv_cache_baseself.kv_pool)returnoutputAgent 场景下如果每次工具返回都做完整 Prefill——一个 5 轮 Agent 调用等于 5 个新请求。增量续接后相当于只有第一次是 Prefill后面 4 次是增量 Prefill Decode。CANN 上 Agent 推理的调度// Agent 场景下的 CANN 推理调度——多 Agent 并行classAgentBatchScheduler{// 同时跑 K 个 Agent每个 Agent 在不同推理阶段std::vectorAgentStateagents;voidScheduleStep(){// 调度策略把处于推理阶段的 Agent 合并成一个 Batchstd::vectorintinferring_agents;for(autoagent:agents){if(agent.phaseAgentPhase::INFERRING){inferring_agents.push_back(agent.id);}}// 合并推理——这些 Agent 都在等模型输出// 它们的 KV Cache 长度可能不同Continuous Batching 策略if(!inferring_agents.empty()){BatchInput batch_input;for(intid:inferring_agents){autoagentagents[id];// 每个 Agent 当前只需 Decode 1 个 Tokenbatch_input.input_ids.push_back(agent.next_input_id);// 每个 Agent 的 KV Cache 位置不同batch_input.kv_start_pos.push_back(agent.kv_cache.current_len);batch_input.kv_lengths.push_back(agent.kv_cache.current_len);}// 一次推理出所有 Agent 的下一个 Tokenstd::vectorintnext_tokensmodel.batch_infer(batch_input);// 分发结果for(size_t i0;iinferring_agents.size();i){intidinferring_agents[i];agents[id].last_outputnext_tokens[i];agents[id].phaseAgentPhase::PARSING;// 切到解析}}// 解析阶段的 Agent——检查是否要调工具for(autoagent:agents){if(agent.phaseAgentPhase::PARSING){agent.CheckToolCall();}}}};Agent 推理的性能瓶颈不在单次推理快慢而在整个循环的吞吐。一个 8 轮 Agent 的端到端延迟 Prefill(80ms) 8 × Decode(100ms) 8 × 工具调用(50ms) ≈ 1.5 秒。CANN 的 Continuous Batching 增量 Prefill 能把多个 Agent 的推理合并8 个 Agent 并行时吞吐从 0.7 agent/s 提到 3.8 agent/s。参考仓库Agent 推理配方GE 图引擎增量执行Runtime 显存管理