本文基于昇腾CANN和昇腾NPU，围绕 cann-recipes-infer 仓库的相关技术展开。
AI Agent 的推理模式跟普通对话不一样：模型输出的是一个动作序列------调工具、看结果、再调工具，直到任务完成。这对推理系统的要求从"一次预测"变成了"循环直到停止"。CANN 在 Agent 推理场景下需要做两件事：高吞吐的循环推理 + 工具调用的 Host-Device 交互。
Agent 推理循环的结构
python
# AI Agent 的典型推理循环------工具调用直到完成
class AgentExecutor:
"""
一个 Agent 推理单元：模型 + 工具集合 + 停止条件
"""
def __init__(self, model, tools):
self.model = model # LLM 模型
self.tools = {
"search": ToolSearch(),
"calc": ToolCalculator(),
"code": ToolCodeRunner(),
"read": ToolFileReader(),
}
self.max_turns = 10
self.history = []
def run(self, user_query):
"""
一次用户请求 = N 轮模型推理 + M 次工具调用
"""
messages = [{"role": "user", "content": user_query}]
for turn in range(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)
if not tool_calls:
# 没有工具调用 → Agent 完成了
return response
# Step 3: 执行工具调用------Host 侧
for tc in tool_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
python
# Agent 推理的关键优化：KV Cache 续接
class AgentKVCacheManager:
"""
管理 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 = 0
def extend_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 # 告诉模型从哪个位置开始
)
# 写到预先分配的位置
for layer in range(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_pos
return start_pos
def incremental_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
)
return outputAgent 场景下，如果每次工具返回都做完整 Prefill------一个 5 轮 Agent 调用等于 5 个新请求。增量续接后相当于只有第一次是 Prefill，后面 4 次是增量 Prefill + Decode。
CANN 上 Agent 推理的调度
cpp
// Agent 场景下的 CANN 推理调度------多 Agent 并行
class AgentBatchScheduler {
// 同时跑 K 个 Agent，每个 Agent 在不同推理阶段
std::vector<AgentState> agents;
void ScheduleStep() {
// 调度策略：把处于"推理阶段"的 Agent 合并成一个 Batch
std::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 (int id : inferring_agents) {
auto& agent = agents[id];
// 每个 Agent 当前只需 Decode 1 个 Token
batch_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 的下一个 Token
std::vector<int> next_tokens = model.batch_infer(batch_input);
// 分发结果
for (size_t i = 0; i < inferring_agents.size(); i++) {
int id = 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。