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

AI Agent 推理：从单次对话到多轮工具调用

news 2026/5/23 19:17:55

本文基于昇腾CANN和昇腾NPU，围绕 cann-recipes-infer 仓库的相关技术展开。

AI Agent 的推理模式跟普通对话不一样：模型输出的是一个动作序列——调工具、看结果、再调工具，直到任务完成。这对推理系统的要求从"一次预测"变成了"循环直到停止"。CANN 在 Agent 推理场景下需要做两件事：高吞吐的循环推理 + 工具调用的 Host-Device 交互。

Agent 推理循环的结构

# AI Agent 的典型推理循环——工具调用直到完成classAgentExecutor:""" 一个 Agent 推理单元：模型 + 工具集合 + 停止条件 """def__init__(self,model,tools):self.model=model# LLM 模型self.tools={"search":ToolSearch(),"calc":ToolCalculator(),"code":ToolCodeRunner(),"read":ToolFileReader(),}self.max_turns=10self.history=[]defrun(self,user_query):""" 一次用户请求 = N 轮模型推理 + M 次工具调用 """messages=[{"role":"user","content":user_query}]forturninrange(self.max_turns):# Step 1: 模型推理——生成回复response=self.model.chat(messages)# Step 2: 解析是否要调工具# Agent 模型的输出格式通常是：# <tool_call>{"name": "search", "args": {"query": "..."}}</tool_call>tool_calls=parse_tool_calls(response)ifnottool_calls:# 没有工具调用 → Agent 完成了returnresponse# Step 3: 执行工具调用——Host 侧fortcintool_calls:tool=self.tools[tc["name"]]result=tool.run(tc["args"])# 把工具结果追加给模型messages.append({"role":"tool","tool_call_id":tc["id"],"content":result})# Step 4: 继续下一轮——带着工具结果再次推理return"Agent 达到最大轮数，未完成"

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_seq=65536):# 预分配一个大的 KV Cache 空间self.kv_pool=torch.empty(num_layers,2,max_seq,num_heads,head_dim,dtype=torch.float16,device="npu")self.current_len=0defextend_with_tool_result(self,tool_result_tokens):""" 工具调用结果追加到已有 Cache 后面 流程： Prefill 阶段（首次）：Query → 模型 Cache[0..q_len) 工具调用后：Cache[0..q_len) 不动，新 Cache[q_len..) 追加 """start_pos=self.current_len end_pos=start_pos+len(tool_result_tokens)# 在 CANN 上做"增量 Prefill"# 不是从头算全部，只算新增部分的 K、V# 模型的其他层 Cache 不动new_kv=self.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_len=end_posreturnstart_posdefincremental_prefill(self,tokens,start_pos):""" 增量 Prefill——只算新增 Token 的 K、V CANN Runtime 支持这种"从指定位置开始"的执行模式 不需要重新构建整图的 KV Cache """# CANN 上通过 GE 的 Incremental 执行模式# 设置输入位置偏移，模型自动从 start_pos 开始output=model.execute_with_offset(input_ids=tokens,position_offset=start_pos,kv_cache_base=self.kv_pool)returnoutput

Agent 场景下，如果每次工具返回都做完整 Prefill——一个 5 轮 Agent 调用等于 5 个新请求。增量续接后相当于只有第一次是 Prefill，后面 4 次是增量 Prefill + Decode。

CANN 上 Agent 推理的调度

// Agent 场景下的 CANN 推理调度——多 Agent 并行classAgentBatchScheduler{// 同时跑 K 个 Agent，每个 Agent 在不同推理阶段std::vector<AgentState>agents;voidScheduleStep(){// 调度策略：把处于"推理阶段"的 Agent 合并成一个 Batchstd::vector<int>inferring_agents;for(auto&agent:agents){if(agent.phase==AgentPhase::INFERRING){inferring_agents.push_back(agent.id);}}// 合并推理——这些 Agent 都在"等模型输出"// 它们的 KV Cache 长度可能不同，Continuous Batching 策略if(!inferring_agents.empty()){BatchInput batch_input;for(intid:inferring_agents){auto&agent=agents[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::vector<int>next_tokens=model.batch_infer(batch_input);// 分发结果for(size_t i=0;i<inferring_agents.size();i++){intid=inferring_agents[i];agents[id].last_output=next_tokens[i];agents[id].phase=AgentPhase::PARSING;// 切到解析}}// "解析阶段"的 Agent——检查是否要调工具for(auto&agent:agents){if(agent.phase==AgentPhase::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。