一、背景介绍

大语言模型原生不具备记忆能力，上下文完全依赖输入 Prompt 携带历史对话。在 Agent 智能体多轮交互场景，Short‑term memory 短期记忆负责保存当前会话最近交互信息，维护对话上下文、工具调用中间状态、会话临时变量，会话结束即失效；区别于长期记忆 Long‑term memory（持久化落盘、跨会话读取）。

在昇腾 NPU 推理链路中，短期记忆模块与 KVCache 推理缓存既有联系也有区别：

● KVCache：Transformer 推理算子层面缓存 Key‑Value 张量，加速 token 生成，属于模型底层推理缓存；

● Short‑term memory：应用 / Agent 业务层会话记忆，维护对话消息列表、工具返回、会话元数据，输出组装给模型输入。

实际开发中两者配合：STM 维护业务会话数据，组装成 prompt 送入模型；推理阶段 KVCache 做张量加速推理。当会话轮次变长，STM 需要做窗口裁剪、摘要压缩，规避 Prompt 超限、NPU 显存占用过高问题。

适用场景

1. 多轮对话聊天机器人，保存本轮会话聊天历史；

2. Agent 智能体，记录工具调用、中间思考、函数返回结果；

3. 任务式 AI 应用，保存任务临时状态，会话关闭全部释放。

二、Short‑term memory 开发模式

2.1 三种开发接入模式

开发模式	适用场景	昇思 MindSpore 接入方式	核心特点
原生列表模式（轻量）	简单多轮对话，会话并发少	Python List 维护 messages，推理前直接拼接输入	零依赖；长会话会触发 token 超限，需自行实现滑动窗口
滑动窗口 STM 模式（常用）	Agent、多轮对话生产应用	封装记忆类，实现窗口截断、头部丢弃、简单摘要；对接 MindFormers 推理流水线	限制最大 token 长度，控制输入 Prompt 规模，适配 NPU 显存约束
缓存 + 摘要混合模式	复杂 Agent，中等长度会话	超过窗口阈值自动对早期对话做昇思大模型摘要，摘要 + 最新对话共同组成短期记忆	保留关键信息，降低 token 消耗，增加推理摘要开销

注意：Short‑term memory不做磁盘持久化，进程重启、会话销毁记忆全部清空；需要跨会话记忆请接入 Long‑term memory 向量库检索。

2.2 整体数据流

用户输入 → Short‑term memory读取历史消息 → 合并新消息，执行窗口/摘要处理 → 组装prompt文本
→ MindFormers送入昇腾NPU推理（KVCache增量推理）→ 模型输出结果写入STM会话缓存 → 返回给用户

2.3 关键约束（昇腾 NPU 部署重点）

1. STM 输出拼接后的总 token 长度不能超过模型配置max_seq_length，否则推理报错；

2. 并发会话场景，每路会话独立一份 STM 实例，做好内存回收，避免 NPU 侧 KVCache 与主机侧会话内存泄漏；

3. KVCache 的缓存长度与 STM 输出 token 强相关，npuMemSize参数需要预留足够缓存空间。

三、核心模块内容详解

3.1 短期记忆数据结构

会话单元示例，每一条记忆单元存储角色、内容、时间戳、工具元信息：

memory_item = {
    "role": "user|assistant|tool|system",

    "content": "文本内容",
    "timestamp": 1750000000,
    "tool_meta": {"tool_name": "", "tool_return": ""}
}

短期记忆容器本质是消息队列，新增消息追加队尾；滑动窗口策略从队头丢弃旧消息。

3.2 核心接口定义

1. add_memory(item)：新增一条会话记忆；

2. get_memory()：获取完整短期记忆消息列表；

3. trim()：窗口裁剪，控制总 token 数，超过阈值丢弃早期消息；

4. clear()：清空当前会话短期记忆；

5. build_prompt()：把记忆列表转换为模型可识别输入 prompt 字符串。

3.3 与昇思 MindFormers 推理对接要点

● MindFormers 的TextGenerator接收输入文本，文本由 STM 的build_prompt()输出；

● 推理使用增量推理，开启 StaticCache 静态 KV 缓存，提升多轮生成性能；

● trim 裁剪逻辑必须在每一轮推理前执行，防止序列超长导致 NPU 推理失败。

四、代码实战（昇思 MindSpore + MindFormers）

环境依赖：mindspore‑ascend，mindformers，运行设备昇腾 NPU；示例实现滑动窗口 Short‑term memory。

4.1 Short‑term memory 封装实现

import time
from typing import List, Dict
from mindformers import AutoTokenizer

class ShortTermMemory:
    def __init__(self, tokenizer: AutoTokenizer, max_token_limit:int=1024):
        self.tokenizer = tokenizer
        self.max_token_limit = max_token_limit
        self.memory_list:List[Dict] = []

    def add_memory(self, role:str, content:str, tool_meta:Dict=None):
        item = {
            "role": role,
            "content": content,
            "timestamp": int(time.time()),
            "tool_meta": tool_meta if tool_meta else {}
        }
        self.memory_list.append(item)
        self.trim()

    def get_total_tokens(self) -> int:
        """统计当前记忆总token数"""
        texts = []
        for msg in self.memory_list:
            texts.append(f"{msg['role']}:{msg['content']}")
        tokens = self.tokenizer(texts, return_tensors=None)
        return sum(len(t) for t in tokens["input_ids"])

    def trim(self):
        """滑动窗口裁剪：超过token上限，从头部丢弃旧消息"""
        while self.get_total_tokens() > self.max_token_limit and len(self.memory_list) > 1:
            self.memory_list.pop(0)

    def build_prompt(self) -> str:
        """将记忆列表组装成对话prompt文本，送入大模型"""
        prompt = ""
        for msg in self.memory_list:
            prompt += f"<|{msg['role']}|>\n{msg['content']}\n"

        return prompt

    def clear(self):
        self.memory_list.clear()

4.2 完整调用推理流程

from mindformers import AutoTokenizer, AutoModelForCausalLM

# 1.加载tokenizer与模型（昇腾NPU）
model_name = "MindSpore‑Lab/Qwen‑7B‑Chat‑FP16"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# 2.初始化短期记忆，最大限制1024token
stm = ShortTermMemory(tokenizer=tokenizer, max_token_limit=1024)

# 模拟多轮会话
dialogue_rounds = [
    ("user", "你好，介绍昇思MindSpore"),
    ("assistant", "昇思MindSpore是全场景深度学习框架，适配昇腾NPU"),
    ("user", "怎么实现短期记忆？")
]

for role, content in dialogue_rounds:
    stm.add_memory(role, content)

# 获取组装好的prompt
input_prompt = stm.build_prompt()
print("组装Prompt:\n", input_prompt)

# 推理生成
inputs = tokenizer(input_prompt, return_tensors="ms")
outputs = model.generate(**inputs, max_new_tokens=200)
resp_text = tokenizer.decode(outputs[0], skip_special_tokens=True)

# 将模型输出写回短期记忆
stm.add_memory("assistant", resp_text)
print("\n模型输出：", resp_text)

# 会话结束，清空短期记忆
stm.clear()

4.3 KVCache 静态缓存开启（性能优化）

结合昇思静态缓存 StaticCache 加速多轮推理，减少重复计算，关键片段：

from mindformers.infer.static_cache import StaticCache
import mindspore.ops as ops

past_key_values = StaticCache(
    config=model.config,
    max_batch_size=1,
    max_cache_len=1024,
    dtype=model.dtype
)

每一次stm.clear()重置会话时，需要同步重置 StaticCache 缓存，防止旧会话 KV 残留。

五、常见问题与调优实践

Q1：推理报序列长度超限

● 根因：STM 未做 trim，对话历史累积 token 超过模型max_seq_length；

● 解决：调小max_token_limit；或者开启摘要模式，对早期对话做大模型摘要压缩。

Q2：并发多会话内存上涨

● 根因：会话销毁没有调用stm.clear()，同时 StaticCache 没有重置；

● 解决：会话生命周期结束，同时清空 STM 列表 + 重置 StaticCache 实例。

Q3：短期记忆与 KVCache 混淆

● KVCache：底层张量缓存，算子加速；进程内内存，不保存业务语义；

● Short‑term memory：业务层会话消息，保存 role/content 语义，用于组装 Prompt。

Q4：需要跨会话记忆怎么办？

Short‑term memory 只负责当前会话，跨会话需要接入 Long‑term memory，向量数据库做检索召回，召回内容注入 STM。

六、总结

Short‑term memory 是昇思大模型 Agent 应用会话层基础组件，主要维护当前会话全部临时交互状态。开发时区分三种接入模式，轻量场景直接列表；生产 Agent 优先滑动窗口 STM；长会话使用摘要混合模式。

开发关键点：

1. 每轮对话必须做 token 裁剪 trim，保护昇腾 NPU 推理不超限；

2. STM 业务记忆与底层 KVCache 分开理解，两者协同工作；

3. 会话生命周期管理，会话销毁必须清空记忆与推理缓存，规避内存泄漏。