最近准备各种面试时整理了许多知识，包括Nanobot源码拆解、Agent、AI Infra和LLM知识学习，这里将学习大纲分享出来。

面试时体感是Agent会问得很细、几乎是每项功能/机制都要了解其作用流程（毕竟Agent也就是一个Prompt处理程序）；AI Infra和LLM部分则考到概念部分即可，对注意力架构、推理端Infra框架（vLLM、SGLang）、FlashAttention、国产开源模型框架的理解要求相对深，并且最好要有自己在云端部署开源模型的经验。

1. Nanobot交互模式流程

• bus.publish_inbound(InboundMessage)
  • MessageBus统一对外的消息对象In/OutboundMessage，对内则以OpenAI格式messages处理，是连接Channel与AgentLoop的桥梁
  • MessageBus维护两个asyncio.Queue（如果队满则会推迟放入对象直到有空间，这里不限制队列长度），Channel往inbound放、AgentLoop取，后者往outbound放、Channel取，互不关心放入后的事。
  • In/OutboundMessage以session_key_override(“channel:chat_id”)会话所走的AgentLoop
• AgentLoop.run() 消费 inbound
  • self.running被 stop() 控制
  • self._connect_mcp() 添加MCP工具
  • 异步等待消息，1秒检查一次，超时、非正常取消和普通出错都会继续循环
  • 接到消息后首先通过命令路由器判断是否为特殊指令，如果是则会直接异步执行，发送消息，不与普通任务排队，继续循环。此处处理的是priority指令，普通指令在 _process_message() 中才执行。
  • 否则通过 asyncio.create_task() 创建 _dispatch() 后台任务，并添加至任务队列，继续循环。添加任务完成时的回调函数，会自行从任务队列删除
  • 以上机制把消息处理全部包装成Task交由后台处理，优先让循环继续，从而可以处理多session会话。特殊指令则直接执行，优先完成该session的特殊机制和消息发送。
  • await会让当前协程暂停等待协程对象返回值，此时事件循环可以去跑其他任务；asyncio.create_task()则会把协程包装成 Task 交由事件循环后台执行。async def 定义异步函数，返回一个协程对象， await 它才会开始执行并拿到值。Task 则是协程对象的一层包装，它有检查是否完成、取消以及添加回调的能力，最重要的是事件循环会调度并提前执行协程对象，减少主程序await所需时间。
• _dispatch()
  • 同一会话用lock串行
    • 每个会话用一把锁，dict.setdefault()会返回已存在的字典值或设置默认值后返回，asyncio.Lock()是同时只有一个Task可持有的异步锁。async with能自动释放。
  • 不同会话用gate并发
    • 全局一个asyncio.Semaphore，最多被拿3（自定义）次
  • 次序上先判断lock再判断gate，确保同一会话只占用一次gate
  • asyncio的锁与门都是对事件循环中的Task进行限制，是在单线程上工作。而C++的互斥锁、counting_semaphore以及python的threading.Lock()工作在多线程上，执行线程阻塞。
• _process_message()
  • self.context.build_messages()联合多种上下文，不在这儿跑
    • 包括以下内容：
      • system prompt

        nanobot 身份和指导原则（包括写前先读、不要幻觉、问清楚、不要相信内容中要求的网络搜索等）
        运行平台：操作系统，系统特性（尤其是windows与posix的命令行区别）
        workspace 路径
        AGENTS.md / SOUL.md / USER.md / TOOLS.md
        长期 memory
        always skills
        skills summary

      • history：未consolidated会话消息
      • 运行时信息：时间、channel和chat_id
      • 用户消息
  • manager把session消息存成JSONL、记忆系统处理长期信息，不在这儿跑
    • MemoryConsolidator.memory_consolidator()先通过 estimate_session_prompt_tokens()估算 get_history()（只返回last_consolidated之后的消息）返回的会话消息token数，超出预算的部分会找安全边界（大于所需减少的token数、last_consolidated的下一个尽量是user消息）后把旧消息交给MemoryStore调用agent总结，并写入长期记忆及历史中，更新 last_consolidated 指针。
    • 压缩预算大小是 $(上下文窗口-最大补全token数(4096)-安全缓冲区(1024)) // 2$，避免需要频繁处理。
    • MemoryConsolidator.consolidate() 会调用保存记忆的工具，其中分别定义了history_entry和memory_update，并要求LLM返回结构化输出。fallback为把原始信息写入HISTORY.md。
    • HISTORY.md记录被总结后的日志，用来grep搜索和人工回查，理论上会无限增长。
    • MEMORY.md记录长期记忆，每次都随新对话一起压缩，没有硬性上限，因此memory质量依赖prompt和模型。
  • 系统消息分支，比如subagent、heartbeat、cron会走这条线，除了不会处理command，与普通分支差不多
  • 普通消息分支
    • 先处理普通slash命令（扔给commands.dispatch()）
    • 请求前memory consolidation 压缩记录
    • ContextBuilder.build_messages()
    • _run_agent_loop()
    • _save_turn()把新消息放回sessions管理，sessions.save()保存到文件
    • 后台重memory consolidation 清理记录
• _run_agent_loop()
  • _trim_history_for_budget防止调用工具的循环发生超窗口，这里不处理记忆，直接从最早的历史记录开始逐条删直到budget以内，还会把有意义的对话前的孤儿工具调用历史删掉。
  • chat
    • 各供应商分别使用openai、openai chat compatible、azure openai和anthropic规范的响应格式
    • Provider类实际是做封装统一，只需要处理上述几类即可，主循环只能看到相同类别的LLMResponse和ToolCallRequest，
    • 工具信息和JSON Schema（如有MCP工具也在这）是通过OpenAI API直接作为参数传递的。
    • chat_stream()等待LLM的流式输出，每等到一次响应就立即发送文本，工具调用在响应结束后一起执行
  • tool call and execution
    • 仰赖模型的结构化输出和OpenAI的API封装
    • 核心是arguments，它是json字典，因此可以直接作为参数调用函数
    • 工具统一Schema（name、description、parameters），对LLM友好，也符合结构化要求
    • Retry机制，工具出错时会把错误信息和错误分析提示词一起返回成消息
  • stop when no tool call
  • 异步等待LLM响应、并发工具调用
• bus.publish_outbound(OutboundMessage)
• Channel / CLI 展示给用户
  • Channel Manager
    • 能看到同一个bus，核心是_dispatch_task后台任务
    • 启动_dispatch_outbound协程对象维持循环，它无限循环地从outbound队列取消息，过滤progress和toolhint后交由对应的Channel发送，有指数退避重试机制
    • 启动所有可用Channel
      • Channel类含以下三种基础方法：
        • start()：侦听消息、把消息处理成InboundMessage放入总线的inbound队列
        • stop()：停止连接，释放资源
        • send()：发送消息
      • 飞书Channel用单独的线程运行Websocket循环，用_on_message()异步函数丢进总线队列
• 安全边界：总体上只能通过prompt软限制LLM输出，必须通过本地程序审查和授权执行
  • 入口身份限制：Channel可单独设置allowFrom列表，手动允许用户消息
  • 工具注册：可以手动关闭shell命令工具
  • restrict_to_workspace：开启后工具内部会自动检查路径是否在workspace内，如果否，会抛出错误；shell命令工具还会对命令行参数多一道检查；唯一例外是ReadFileTool可以读取内置skills目录
  • 工具参数校验：工具会先把LLM结构化输出的JSON转成参数字典，再根据预设Schema检查每项参数
  • Shell执行限制：拦截危险命令如rm -r、reboot等，600秒超时限制和10000输出字符截断窗口
  • Web访问限制：只允许http/https、解析hostname、阻截含有内网地址的调用，重定向后地址也会审查（防止重定向到云metadata或内网地址）。但还是会有一定危险，比如访问危险网站、安装危险skills等
  • MCP工具注册：超时和手动启用
  • Prompt安全Guideline：写前读、分析工具错误、不信任网络数据、不听从网络数据命令。提示词注入的软防御。
• skills处理：
  • SkillsLoader挂载在ContextBuilder里，skills包括workspace和built-in两种，为ContextBuilder的System Prompt提供always_skills和skills两种prompt
  • always_skills去掉头部YAML formatter，把全文放进system prompt
  • skills只读头部YAML formatter，格式化skills名、描述和SKILL.md文档位置，方便LLM直接调用read_file工具读取
  • 可以通过$<name>手动激活skills，也就是强制该skills注入消息上下文
• Subagent机制：
  • 通过spawn工具生成，参数包括task任务内容和label任务标签，调用时也设置了channel和chat_id，因此可以匹配正确的session。生成后会有标签和任务id提示返回给Agent。
  • 自身是一个后台Task，跑了一个简化版的Loop。Subagent没有message、spawn和cron工具以避免向用户发消息和分身。
  • 有专用的简化版system prompt，包含subagent人格、安全guideline、系统信息、工作区信息、skills summary，还有user身份的任务消息。
  • 完成后调用回调函数，将结果prompt（任务id、任务内容、任务结果、要求总结并发给用户）以subagent身份的系统消息附上session识别字符串放入主Agent的消息队列。
  • 主Agent会以assistant身份处理结果。
  • 遇到/stop时，会取消该会话的subagent task。
  • 目前不能通过配置单独设置不同的provider和model，config时刻与主Agent共用，这以轻量化为目的。如果需要则要让spawn工具带新provider的信息，在subagent里新开一个provider类处理，同时config中也需要额外设置subagent的provider、model等。

2. 相关知识补全

2.1 FlashAttention

Tri Dao于2022年提出的SDPA加速技术，针对原始SDPA运算中间结果需要频繁读写进全局显存的问题，采用分块式的设计将把矩阵块的在SRAM里尽量算完后写回，创新式地分段计算SDPA，尤其是渐进式更新sigmoid函数，从而逐步更新结果矩阵，缓解了显存读写瓶颈。由于没有存储中间结果，在反向传播时实际上需要重计算中间结果。此外还利用了safe sigmoid维持数值稳定性，主要是所有项先减去行最大值再计算自然指数，还分开计算分子分母，当然也是分块计算的。行最大值和分母被存储，用于计算反向传播的中间结果（中间结果未保存，需要重新计算）。事实上把显存复杂度从平方降低到了线性，读写小矩阵提高了IO速度，综合提供了2-4x的计算加速。这种设计思想叫做IO-Aware。

2.2 Agent框架

2.2.1 MCP vs Skills

MCP是Anthropic推动的开放网络通信协议标准，主要采用服务器-客户端架构。服务器端是一个动态运行的实时服务，为客户端（Agent）提供开放的工具接口和数据。相当于云端工具，由于MCP跟工具一样，它会拥有极详尽的工具定义和参数JSON Schema，并且与工具一起预加载到消息上下文中，因此上下文消耗极高。通过延迟加载/语义搜索发现定向加载MCP工具可以缓解上下文压力。

Skills是开源高级指令规范（Markdown/YAML），包含轻量级指导文本、代码片段、触发条件和工作流步骤等，有独特的按需加载机制。更像是某项业务的指导书，用较短的文本教Agent如何完成某项任务，并且可以使用所提供的功能性代码。

2.2.2 超长上下文管理

2.2.2.1 多模态上下文管理

比如视觉数据转换成token的数量（取决于Patch大小），100万token可能只相当于1000帧、30秒视频。

1. 动态分辨率：高分辨率图像切分成VE能处理的块大小、整体resize成相同块大小，两边每块分别编码并拼接作为视觉特征。本质是降低Patch数量的同时通过不同的切分手段保持一定有效性。
2. mRAG：最原始的做法是把数据通过编码器（如CLIP）压缩成高维向量存入向量库，以RAG的方式进行。但也有使用专门视频、音频标注模型描述每帧画面后再作为信息补充给LLM的。

2.2.2.2 长思维链推理

思维链(CoT)是Google于2022年发现的LLM涌现出的能力，最初分为少样本CoT（给QA及推理过程示例）和零样本CoT（无样本的推理命令，如Let's think step by step）。这过程中可以使用多轮对话，Prompt-Chaining让人手动把问题的解决过程一步步拆开、前一步的结果放进后一步作为示例上下文，但针对任务不同、很费人工，Auto-CoT技术先对问题集聚类、让LLM对于聚类中心的问题自己生成零样本示例再进行少样本CoT。

对于长推理任务，Shunyu Yao于2023年提出了思维树（ToT），定义任务所需思维/步骤的数量以及每步的候选项数，每步都要求评估当前是否有希望得到正确答案，每步选择5个最优的勾选项继续步骤，总体类似于遗传算法。其中评估可以让LLM自己评估，也可以通过强化学习等训练出类似控制器的组件，不过后者比较依赖数据集。

现代Agent通过与外部环境交互可以验证自己的部分推测，我认为这是Claude Code与Codex的/goal模式核心，即一个不变的上下文：要干什么、如何检验、结束条件是什么，结束条件通常是某项指标达标或者一份报告，抑或是消耗token budget。而Karpathy的Autoresearch则想法大致类似，没有明确的结束条件，但是对于实验步骤、可改动文件的规定等的更死，并且主动要求使用git和日志管理进度（主动回滚）。前者是一个根植于Agent程序的上下文构建和循环机制，而后者类似一个明确规定了实验操作的Skills。

AI-Scientist-v2是2026年5月发表在Nature正刊上的科研Agent，比/goal更严格定义科研任务、比Autoresearch更全面地涵盖科研步骤，它进行了科研流程的每一步：对人类提供的话题提出论文idea（包括论文所需要的所有部分）、简单查重、在有限步骤内树形搜索、消融实验、绘图、引用、写文、自动预审稿。然而仍有许多问题：提出幼稚的idea，陷于错误的implementation，幻觉引用等等。最大贡献在于把实验过程的树分为多个阶段，每个阶段有自己专属的节点prompt以区分当前正在做的步骤，并且让Agent自己判断是否进入下个阶段。

但是Agent仍不具备体系化的专业知识（或许可以把某领域的专业知识写成树状的文件结构，再变成skills加上脚本注明retrieve方法，方便agent检索），以及人类般灵活调整目标的能力（这一点则可能要求agent对于目标有回退的能力，即一级目标、二级目标、三级目标这样先提出来，当目标重复实现多次都失败后就记录失败目标、回退至上级目标、要求找其他方法）。或许这些可以成为长任务Agent的下一个目标。

2.2.2.3 记忆管理系统

将单纯的输出结果堆叠进上下文窗口，改造成一个“仿生记忆分层操作系统”，LLM当作CPU、上下文窗口当作RAM，外部存储文件当作硬盘。Nanobot的记忆和上下文系统就是一个例子。

1. 工作记忆：上下文窗口中的内容，包括系统信息、工作区信息、当前任务prompt、长期记忆、最近几轮的对话等等。
2. 情节记忆：过去所有的对话、所有运行过的工具历史，按照时间线写入外部数据库，方便查询。
3. 语义记忆：长期记忆，总结过去一段时间的对话写入外部数据库，被总结的对话将来不再放入上下文。类似于人类的记忆与遗忘。

2.3 一些早于Agent的LLM应用框架

2.2.2 ReAct (Resoning and Act)

Shunyu Yao于2022年提出的LLM提示技术，它将“推理”与“行为”协同，让LLM执行“思考”、“行动”、“观察”的循环思维链，其中“行动”是与外部环境交互的行为（比如搜索），“观察”则是对外部环境的观察（比如搜索结果）。

2.2.3 RAG (Retrieval Augmented Generation)

Meta于2021年提出的LLM提示技术，它提供了一种知识检索的轻量型工具。该工具把大量文本分块并向量化作为知识库，对用户问题中的关键词通过相似度匹配获取最相近的知识块，作为回答用户问题的上下文。

2.4 AI Infra

先浅尝一下大模型的恐怖之处：

FP16精度的500B参数权重差不多刚好需要1000GB即1TB显存

显存墙：每轮训练时反向传播的梯度值大概与模型参数量等同，优化器状态值（如Adam的动量）大概需要4-6倍模型参数量，保存用于计算梯度的前向传播的中间值则有多少要存多少。因此训练时的显存压力至少是推理时的6倍，这还需要乘上Batch Size。

算力墙：训练时反向传播的计算量大约是前向传播的2倍，因此训练时的算力需求至少是推理时的3倍。

此外训练还需要同构集群，分布式训练需要高带宽的InfiniBand网络（通信墙），以及checkpointing容错机制处理某些硬件宕机的情况。

2.4.1 训练端：分布式并行策略

以下来自[Megatron-LM]

2.4.1.1 GPU间通信机制 (NCCL Primitives)

1. Broadcast：单卡对所有其他卡广播。
2. Gather与Scatter：分别是让所有卡发数据给单卡拼起来、单卡把数据切碎发给所有卡。
3. All-Gather：所有卡广播自己的数据，并把收到的数据拼起来。
4. Reduce：规约，让所有卡把数据发给单卡，在接收过程中还能进行数学计算，通常是相加，单卡直接拿到数据之和。
5. All-Reduce：全规约，让所有卡把数据相加，并把数据之和广播让所有卡都拿到。
6. Reduce-Scatter：规约分发，让所有卡先把数据相加，再把数据之和切碎并分发给所有卡。

2.4.1.2 Data Parallelism (ZeRO)

传统做法DPP：是每张卡上都有完整的模型权重，前向传播时每卡输入不同的Batch数据，各自反向传播算出梯度后，通过All-Reduce对所有卡计算得的梯度求平均，最后各自更新自己的优化器状态。这样每张卡都必须装下完整模型权重、梯度和优化器状态，并且每张卡的模型权重实际相近，非常冗余且对大模型不可行。

1. ZeRO-1：优化器状态切分，每卡有完整模型和梯度，但优化器状态在所有卡上均分（即每张卡负责一部分参数的更新），每张卡计算完梯度后需要一次Reduce-Scatter平均并分发所有梯度，更新完自己负责的参数后，还需要一次All-Gather让所有卡拿到重新拼好的参数。均分优化器状态已经解决了最大的显存瓶颈，由于All-Reduce=Reduce-Scatter+All-Gather，因此相比DPP没有增加通信量。
2. ZeRO-2：梯度切分，进一步切分梯度，每卡只要计算局部梯度，通过Reduce传给这部分梯度对应的卡。更新完成后还是需要All-Gather拼参数，但依旧没有增加通信量。
3. ZeRO-3：参数切分，最后切分模型参数，前向计算通过All-Gather临时把参数拼起来给所有卡，算完立即释放。更加依赖高频通信，相比DPP真正提高了通信量。

2.4.1.3 Tensor Parallelism

单层网络的参数量单卡已经装不下时，就需要张量并行（层内并行），它对自注意力层和MLP进行了矩阵切分。以MLP为例，它将矩阵一分为二：

1. 假如我们有$Z=GeLU(XW_1)W_2$

   $$X\in R^{B \times D}, W_1\in R^{D \times B}, W_2\in R^{B \times D}$$ $$Y=GeLU(XW_1) \in R^{B \times B},\; Z=YW_2\in R^{B\times D}$$

2. 先把$W_1$按列切分分别放到两张卡上，$X$复制到两张卡上，各自计算一部分输出；再把$W_2$按行切分分别放到两张卡上，计算第二次矩阵乘，这次输出必须通过通信相加才能得到最终输出。

   $$W_{1,1},W_{1,2}\in R^{D\times \frac{B}{2}},\; W_{2,1},W_{2,2}\in R^{\frac{B}{2} \times D}$$ $$Y_1=GeLU(XW_{1,1}),\;Y_2=GeLU(XW_{1,2}) \in R^{B \times \frac{B}{2}}, \\ Z_1=Y_1 W_{2,1},Y_2 W_{2,2}\in R^{B\times D},\;Z=Z_1+Z_2$$
3. 因此用TP的每层都需要所有卡的通信，由于其频繁性，TP通常只能在同机器内通过NVLink运行。

在多头注意力中则可以按注意力头切分，即按行切分投影矩阵，让不同GPU负责不同头，计算完SDPA后乘以自己那部分对应的输出投影矩阵，最后All-Reduce相加即得到输出。

2.4.1.4 Pipeline Parallelism

TP每层通信对于极深的模型可能造成较大时延，我们可以进行层间切分（流水线并行），比如8层的模型让GPU 0 负责1-4层、GPU 1 负责5-8层，数据变得流水线传递。

单纯流水线的痛点是每轮GPU 0 把输出交给GPU 1后，直到后者把反向传播梯度传回来前都闲置着（流水线气泡）。

1F1B (One Forward, One Backward)策略把大Batch输入拆成小Micro-batch，算完一个Mb的前向后可以立马切换去算前一个Mb的反向。

2.4.1.5 其他问题

1. TP开大后训练速度为什么可能变慢，通信开销来自哪？
   • 通信开销是每个被切分的层都要执行一次All-Reduce把输出拼出来。
   • 变慢的可能原因有：
     • 并行数量超出单机卡数，从而TP开始依赖外部通信网络，外部InfiniBand带宽远低于内部NVLink。因此TP并行数不要超过单机卡数。
     • 矩阵过小导致没喂饱GPU，无法发挥Tensor Core的峰值算力。
2. 训练时具体如何选择分布式策略？
   • TP：单层参数量极大时，优先设为单机GPU数量。
   • PP：看单机TP能装下几层，至少设为能装下所有层的机器数量；Micro-Batch数量则需要额外调整，通常是PP的4-8倍。
   • DP：TP * PP即一个模型副本所需要的最小GPU数量，集群里剩下的所有卡可以用于DP以尽量增大Batch Size。再加上ZeRO-1减少优化器显存。
3. Fully Async（完全异步）训练会带来什么问题？
   • Fully Async（异步SGD）指算完梯度的节点不等待其他节点平均梯度，直接去更新全局参数
   • 这会导致算得慢的节点用来自版本号较低的模型算出的梯度，强行去更新版本号已经更高的模型。
4. 训推分离的好处？
   • 训练需要同构集群和InfiniBand；推理可以无需训练级算力、处理好高并发即可
   • 训练是一个长时间的批处理任务，不能接受任何外部干扰；推理的并发量波动较大，与训练集群可能抢占显存或带宽资源，所以应当保持隔离。

2.4.2 推理端：长上下文推理优化

推理端的核心痛点在于，自回归解码的范式需要前面所有Token的KV矩阵，为了不重复计算，它们会被存进显存中，即KV Cache。它的体积可能巨大，而每个用户由于请求token数不同，每块KV Cache的体积差异很大，不适合分配同样大小连续的显存空间，即显存碎片化问题，包括请求了没用上的内部碎片和空隙不够用的外部碎片。

2.4.2.1 vLLM与PagedAttention

vLLM是UC伯克利开源的主流LLM高速推理与部署服务框架，最主要特性是PagedAttention和连续批处理。

Woosuk Kwon于2023年提出PagedAttetion，把KV Cache自身切成token数较低（如16）的小块，维护一张逻辑表记录所有小块的物理地址，这样基本消除了显存外部碎片、内部碎片也因为分块体积较小可以忽略不计。思想上与虚拟内存、内存池/对象池的概念比较相近。Batch Size较大时这些频繁访存的开销几乎可以忽略不计，而吞吐量可以翻倍。

它还有几项好处：

1. 连续批处理：由于GPU只能计算矩形矩阵乘法，所以传统批处理都需要把短请求padding到与最长请求一样长，从而导致大量无效计算，而新请求由于显存被占满又进不来。每次批处理计算时都只需要计算固定大小的KV矩阵与请求token，较长的请求和KV Cache被分解成与其他请求大小相同的workload，从而使较短的请求在计算完成后，它的Cache位置可以被标记为可用，从而让新的请求token进来。
2. Prompt Caching (SGLang)：如果token块也被切块，为它计算全局唯一的哈希值，并且配备指向KV Cache地址的指针。新请求进来后也会被切块计算哈希值，如果能匹配到，则可以直接引用已存在的KV Cache。而显存中的token块被引用时会累积引用请求数，并且使用一种特殊的前缀树：基数树（Radix Tree）管理Cache新增与剔除。在基数树中，被引用的节点计数增加，新节点成为该节点的子节点，当节点数量已满时，驱除计数最少的节点。

2.4.2.2 长上下文优化

2.4.2.2.1 RingAttention

当一个用户的长请求Q和KV Cache无法在一张卡中完整保存时，需要将Q和长Cache切分成多个存在不同卡中。

RingAttention利用了FlashAttention的Online Softmax算法，Q_i固定在一张卡上，当前有的KV_i块会与Q_i计算SDPA以更新临时输出O_i，而每轮计算的同时向环状拓扑上的下个GPU发送自己的或者上次收到的KV_j块，而自己也会收到来自上个GPU的KV_k块。当某个KV_i块回到时i点，每张卡的Q_i都已经用所有KV_i块更新过O_i了，也就得到了输出。

StripedAttention瞄准自回归推理的因果性发现了前者的缺点，Q_0不需要KV_1及以后的KV，因此GPU 0有大量时间都在无效运算，而GPU n-1跑的时间最长。因此对Q采用交织化/条带化分配，给所有卡按序轮流分配token，从而几乎让所有卡的负载均衡。

Context Parallelism (CP, Megatron-LM)则更进一步，针对复杂集群提出了优化。首先由于拓扑复杂，P2P Ring通信并不高效，CP在前向计算前用TP的方式All-Gather收发KV块，反向计算后用Reduce-Scatter分布梯度；其次引入双缓冲机制，在计算当前SDPA时，下一个KV块已经在进来了；最后对于所有MLP激活值也进行切分，与Sequence Parallelism的不同点则在于后者只切分了Layernorm和Dropout输出值。

2.4.2.2.2 KV Cache压缩与丢弃

1. H2O (Heavy-Hitter Oracle) / SnapKV 发现Attention权重集中在少数token上，如首句、特殊标点和刚刚提到的物体，因此推理时可以动态评估token重要性，丢弃不重要的KV Cache。以及还有其他各种各样的压缩方式，比如通过相似性、频域信息等判断token重要性，MLLM还要压缩其他数据的token。
2. 低比特量化：常见的是将FP16 KV Cache压成INT8、INT4、FP8。

2.4.2.3 投机采样

每次生成一个token使得Tensor Core算力相对大模型从全局显存读取权重的时间被浪费了，因此可以引入一个小模型（Draft Model，比如1B）先猜出下面5个token。接着5个token一起被喂给大模型（如70B），让它通过一次前向计算验证这些token的合理性。猜对的token和大模型新生成的token一起返回，从而让token吞吐量翻倍。

2.5 现代LLM框架

2.5.1 LLM的新组件

2.5.1.1 RoPE

旋转位置编码源自Jianlin Su于2021年发表的RoFormer论文，这是一种可以表示相对位置关系的乘性绝对位置编码，具有数值稳定性、易外推性和极低的计算复杂度。让我们从二维开始：

假设我们有一个位置索引为m的2维query向量$q$，一个位置索引为n的2维key向量$k$。我们需要它们经过位置编码后的内积能是m-n的函数，因此自然联想到复数表示，并且我们用复数表示向量，即

$$f(q, m)=q e^{im\theta}=(q_1 + i q_2)(\cos im\theta + i\sin im\theta)=\begin{bmatrix}&\cos m\theta &-\sin m\theta \\ &\sin m\theta &\cos m\theta\end{bmatrix} \\ <f(q,m), f(k,n)>= qk^*e^{i(m-n)\theta}$$

作者把高维向量分成$d/2$个子空间，每个空间上进行二维旋转，就得到了高维向量的旋转编码，可以直观理解为每个平面是向量关于某个轴进行旋转时的旋转平面。此外，为了捕捉不同频域的信息，每个子空间i的$\theta_i=10000^{-(2i-1)/d}$。由于每维结果都相当于两个相邻分量加权加减，因此高维向量的旋转编码结果还可以通过简单的逐元素相乘和变换元素位置获取。

它的易外推性体现在通过改变$\theta$的大小可以无需训练或少训练地让可容纳绝对位置变大，新的位置相对旧位置相当于内插结果，对于LLM来说易于理解。

2.5.1.2 激活函数和门控单元

We offer no explanation as to why these architectures seem to work; we attribute their success, as all else, to divine benevolence. —— Noam Shazeer, GLU Variants Improve Transformer, 2020

下文中$\sigma(x)=sigmoid(x)$。

1. GLU (Gated Linear Units)于2017年提出，来自LSTM的门控机制，作为一种神经网络层，通用形式为$GLU(x,W,V,b,c)=\sigma(Wx+b) \odot (Vx+c)$，相当于线性层接门控机制，通常认为用更多的参数量换到更好的效果；
2. GELU (Gaussian Error Linear Unit)于2016年提出，$GELU(x)=x \odot \Phi(x)\;\Phi$是标准正态分布的累积误差函数，优点是导数连续、不会剪裁梯度、近似计算的复杂度小；
3. GEGLU是把门控单元替换为了GELU的GLU，$GEGLU(x,W,V,b,c)=GELU(Wx+b) \odot (Vx+c)$；
4. SwiGLU
   1. Swish由Google Brain于2017年提出，$Swish(x)=x \odot \sigma(x)$，优点是非线性、不会梯度消失，缺点是sigmoid计算量；
   2. SwiGLU于2021年由Google DeepMind提出，采用Swish作为激活函数，$SwiGLU(x,W,V,b,c)=swish(Wx+b)\odot (Vx+c)$；

效果最好的是GEGLU和SwiGLU。

2.5.1.3 归一化

归一化的意义：稳定化梯度流（避免深层网络中梯度爆炸/消失）、防止输入因内部协变量漂移（激活值逐渐偏移）。

为什么BatchNorm失效：文本请求的长度参差不齐；训练时Pipeline Parallelism的Batch Size较小，BN的均值及方差变化很大。

1. LayerNorm于2016年提出，不在Batch维度，而是在特征维度上逐token进行归一化：

   $$\mu = \frac{1}{H}\sum\nolimits^H_{i=1}x_i \\ \sigma^2=\frac{1}{H}\sum\nolimits^H_{i=1}(x_i-\mu)^2 \\ y=\frac{x-\mu}{\sqrt{\sigma^2+\epsilon}}\odot \gamma + \beta$$

   $\gamma,\beta$为可学习参数，$\epsilon$为防止除零的保护值。

2. RMSNorm (Root Mean Square Layer Normalization)于2019年提出，动机是发现LN的有效性主要来源于梯度稳定，即只有方差缩放有用、平移几乎无用，因此删去，还能带来少一次遍历的性能提升。

   $$RMS(x)=\sqrt{\frac{1}{H}\sum\nolimits^H_{i=1}x_i^2+\epsilon} \\ y = \frac{x}{RMS(x)}\odot \gamma$$

3. Pre-Norm vs Post-Norm：

   • Post-Norm是Attention is All You Need原文使用的归一化位置，即在残差链接之后归一化$y=Norm(x+F(x))$，使得残差连接变得不纯粹，反而容易梯度消失（深度越大、归一化分母变得指数级增长，深层反向传播回来的梯度也被缩放多次）需要复杂的学习率预热技巧；
   • Pre-Norm把归一化位置放在输入层前，即$y=x+F(Norm(x))$，它的梯度不会被指数衰减，因此更容易学习。不过这样容易导致最后一层前的输入x变得方差很大，因此预测层前也需要归一化；
   • 有研究发现Pre-Norm的训练最终效果可能不如Post-Norm，Pre-Norm的输出相比Post-Norm不太依赖残差项，结合深层展开后的输出表达式，可能说明Pre-Norm残差连接更接近在同一层中平铺而非加深层数。

2.5.1.4 注意力机制

在本文中，MLA被视为GQA的一般化，它用投影矩阵的方式替代了GQA的分割、重复，并引入了一个恒等变换技巧来可以进一步压缩KV Cache，同时采用了一种混合方法来兼容RoPE。总的来说，MLA称得上是一种非常实用的注意力变体。 —— Jianlin Su, 缓存与效果的极限拉扯：从MHA、MQA、GQA到MLA，2024

从最初的MHA开始，主流注意力机制的演变执着于降低KV Cache的体积同时保持较好的推理性能。

1. Multi-Head Attention (MHA)来自Attention is All You Need，在特征维度上切分，相当于多个独立单头注意力拼接： $$x_1, x_2, ..., x_l \in \mathcal{R}^d \\ \{q,k,v\}_i^{(s)}=x_iW_{\{q,k,v\}}^{(s)}\in\mathcal{R}^{d_{\{k,v\}}},\; W_{\{q,k,v\}}\in\mathcal{R}^{d\times d_{\{k,v\}}} \\ d_k=d_v=d/h$$
2. Multi-Query Attention (MQA)来自2019年的论文，直接让所有头共享同一个KV，从而将KV Cache能减少到$1/h$，并且损失不大，参数量可以通过FFN/GLU弥补。
3. Grouped-Query Attention (GQA)来自2023年的论文，把所有头分成$h \% g = 0,\;g$个组，每组共享一对KV，从而将KV Cache减少到$g/h$，从而提供了MHA到MQA和自然过渡： $$\{k,v\}_i^{(\lceil sg/h\rceil)}=x_iW_{\{k,v\}}^{(\lceil sg/h\rceil)}\in \mathcal{R}^d_{\{k,v\}},\; W_{\{k,v\}}^{(\lceil sg/h\rceil)}\in \mathcal{R}^{d\times d_{\{d,v\}}}$$ LLAMA2/3-70B中$g=8$，因为单机上通常只能装8卡，从而尽量增大组数同时减少卡间通信。

4. Multi-Head Latent Attention (MLA)来自DeepSeek-V2，希望通过投影矩阵提升GQA的能力，具体是先把$x_i$投影成任意维的$c_i$，再用$c_i$投影获得$k_i,v_i$。根据下面的推导，我们只需要存$c_i$作为Cache，当$c_i$维度为$g(d_k+d_v)$时Cache的大小就与GQA相同，但不需要更强的能力只为了减少Cache时可以降维（DeepSeek-V2为512）：

   $$c_i=x_iW_c\in\mathcal{R}^{d_c},W_c\in\mathcal{R}^{d\times d_c} \\ \{k,v\}_i^{(s)}=c_iW_{\{k,v\}}\in\mathcal{R}^{d_{\{k,v\}}},W^{(s)}_{\{k,v\}} \\ q_t^{(s)}k_i^{(s)T}=x_t(W_q^{(s)}W_k^{(s)T})c_i^T$$

   上面把$W_qW_k^T$合并从而存$c_i$代替$k_i$，至于$v_i$则将矩阵吸收进$W_o$即可，这就是理想情况下该方法的优越之处。然而上面是没考虑RoPE的，如果加了RoPE，$W_qR_{m-n}W_k$就没法固定了。因此为了保留RoPE，每个注意力头的Q、K新增$d_r$个维度来添加RoPE，即

   $$\{q,k\}_i^{(s)}=[\{x_i,c_i\}W_{\{q,k\}c}^{(s)},x_iW_{\{q,k\}r}^{(s)}R_i]\in\mathcal{R}^{d_k+d_r},\;\\W_{\{q,k\}c}^{(s)}\in\mathcal{R}^{d\times d_k},W_{qr}^{(s)}\in\mathcal{d\times d_r},W_{kr}\in\mathcal{R}^{d\times d_r}$$

   新增的$d_r=d_k/2=64$只有K Cache。此外为了减少参数量和梯度，Q也先经过一步低秩投影，即

   $$c_i'=x_iW_{c'}\in\mathcal{R}^{d_{c'}},W_{c'}\in\mathcal{d\times d_{c'}}$$

   后面的Q投影矩阵也改为C’投影矩阵，$d_{c’}=1536$。可以发现与MHA相比只是多了一步低秩投影以及只在部分新增的维度加RoPE，虽然增加了计算量，但显著减少了KV Cache，因此推理速度提升效果明显。此外，由于存的是$c_i$而非投影后的KV矩阵，增加注意力头的数量$h$并不会增加KV Cache，因此DeepSeek-V2的$h$大胆地设为128、一般开源模型的3倍。

   总的来说MLA相当于GQA的一种变换一般化版本，整条路径是朝着性能尽量不变、KV Cache越来越少的方向走的。

2.5.2 MoE

部分理论分析参考自Jianlin Su的MoE环游记系列博客。

MoE最早可追溯到Jordan等人于1991年的论文 Adaptive Mixtures of Local Experts，如今LLM的稀疏门控MoE (Sparsely-Gated MoE)则来自Noam Shazeer等人于2017年的论文 Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer，Google则于2021年在 Switch Transformer 中证明了MoE在Transformer中的潜力，MistralAI在2024年的 Mixtral of Experts 中对MoE予以推广。

MoE的影响主要体现在参数量与计算量的解耦，此后模型参数量可以堆至千B以提高知识容量，但推理时只激活极少专家，如13B，使得推理速度极快。

2.5.2.1 MoE机制

MoE层是代替前馈神经网络层（FFN）即线性层的特殊层，主要包含Router与Experts，前者为稀疏门控，后者为许多个较小的FFN，并不改变注意力层。它通过一个路由机制，计算Token对所有专家的权重，最简单的路由机制来自 Mixtral of Experts：

$$G(x)=Softmax(TopK(W_g\cdot x))$$

$W_g$是Router的权重矩阵，假如选择Top-2个最大的概率值，我们就会让这两个Experts根据$G(x)$混合，而未被选中的Experts则不被计算，从而达到一种参数利用的稀疏性。

2.5.2.2 负载均衡

MoE训练的痛点是马太效应，在训练初期如果某些专家被频繁更新，由于它们的表现更好，就会导致路由器会在未来更倾向于把Token交给这些专家，其余专家就会被“饿死”。

为此许多工作设计了各种各样用于负载均衡的辅助损失（Auxiliary Loss），比如让目标分布为均匀分布、用损失要求分配分布接近目标分布，即

$$\mathcal{L}_{aux}=\frac{1}{2}\|F-Q\|^2=\frac{1}{2}\sum\nolimits_{i=1}^{n}(F_i-1/n)^2$$

DeepSeek提出过 Loss-Free 方案，它不改变Router打分结果，而是改变$argtop_k G(x)$这个方式，引入一个可学习的偏置项影响分配（$argtop_k G(x)+\beta$），但不参与权重大小的计算，从而隔离了负载均衡对于LM Loss的负面影响。这里引入偏置项的目标依旧是均匀分布，不同的是，通过损失梯度的推导可以直接获得偏置项的梯度更新方法，这里引用Jianlin Su的推导

$$\mathcal{L_{aux}}=\frac{1}{2}\|\beta+sg[F-\beta]-Q\|^2=\frac{1}{2}\sum\nolimits_{i=1}^{n}(\beta_i+sg[F_i-\beta_i]-1/n)^2 \\ \nabla_\beta\mathcal{L}_{aux}=\frac{1}{2}\nabla_\beta\|\beta+sg[F-\beta]-Q\|^2=F-Q \beta\leftarrow\beta-\gamma(F-Q)\;(SGD) \\ \beta\leftarrow\beta-\gamma sign(F-Q)\;(SignSGD) \\ \beta\leftarrow\beta-\gamma\frac{F-Q}{RMS(F-Q)}\;(RMS\;Adaptive)$$

不过均匀分布也不一定是最好的专家分布，2024年的DeepSeekMoE中还分别提出了不改变总参数量、只提升专家数量的Fine-Grained Expert，以及固定路由某些专家、在剩余专家中选剩余budget的Shared Expert，它们都能提升模型质量，Jianlin Su认为这种人为的不均衡可能模拟了现实世界的非均匀性。

2.5.2.3 专家容量

由于GEMM的计算效率问题，我们希望被使用的专家里，token分配还是相对均衡，因此 Switch Transformer 固定了每个专家可以接收的token数量，即专家容量

$$expert\; capacity = (\frac{tokens\;per\;batch}{number\;of\;experts})\times capacity\;factor$$

capacity factor通常范围为[1, 2]，越大矩阵越大消耗算力越大，越小超出的token数可能越多。少于容量的专家会用Dummy Token补齐以对齐，而超出容量的token则不通过专家FFN计算，直接通过残差连接通往下一层。

2.5.3 长上下文窗口

前面的注意力机制小节已经介绍了MQA、GQA和MLA三种减少KV Cache的方法，它们由于延续了MHA的完整SDPA机制尤其是Softmax，这导致$QK^T$计算不可免，序列计算复杂度总是平方，因此被称为Full Atention。

下面的论文都致力于降低序列计算的复杂度，Linear Attention无心插柳、Mamba有意为之，两者都利用RNN的$O(1)$计算复杂度来改进或替代注意力层，但都具有相似的问题。

2.5.3.1 Linear Attention

Linear Attention于2020年发表，它的Trick在于我们不使用Softmax，改变数学公式，使用一些可以拆成两部分的核函数$\phi(x)$（文中$\phi(x)=elu(x)+1$，更早的Efficient Attention则直接用softmax，妙在可以自动归一化）衡量相似性，从而$\phi(K)^TV$可以被预计算并缓存，$\phi(Q)$只要和这个固定大小的矩阵相乘。因此计算复杂度和显存占用降为线性。

正如原文标题所言，线性注意力层可以规范化成RNN表示，即

$$s_i=s_{i-1}+\phi(x_iW_K)(x_iW_V)^T \\ z_i=z_{i-1}+\phi(x_iW_K) \\ y_i=f_l(\frac{\phi(x_iW_Q)^Ts_i}{\phi(x_iW_Q)^Tz_i}+x_i)$$

不过Jianlin Su也有提到Attention自身的二次性算力增长可能需要序列长度以K为单位时才会体现，在此之前的复杂度主要在FFN层，因此仅对注意力层的改进带来的速度增益可能体感较少。此外线性注意力还有问题是每decode一个token，就相当于新$K^TV$与原来的做加法，位置越后的token见到的$\sum\nolimits_{j=1}^t K^TV$就越大，从而导致之前每个token的作用被均摊，因此后续有其他工作引入遗忘、TTT等新机制解决该问题，可以看blog。

2.5.3.2 Mamba

Mamba的思想来源于RNN，更准确地说来自状态空间模型（SSM）或其后继结构状态空间序列模型（S4）。SSM定义了一个线性ODE系统（A、B、C）：

$$h'(t)=Ah(t)+Bx(t),\; y(t)=Ch(t)$$

其中$h(t),\;x(t)$分别是状态和输入。S4加上一个量化步长$\Delta$，把状态在时间上量化，从而形成了类似线性RNN的状态转移：

$$\bar{A}=\exp(\Delta A),\;\bar{B}=(\Delta A)^{-1}(\exp(\Delta A)-I)\cdot \Delta B\;(ZOH\; descretization \;rule) \\ h_t = \bar{A}h_{t-1}+\bar{B}x_t,\; y_t=Ch_t$$

多步状态转移输出可以拆解成输入经过指数数列卷积，即

$$\bar{K}=(C\bar{B},C\overline{AB},...,C\bar{A}^k\bar{B},...),\;y=x * \bar{K}$$

这个线性系统很重要的假设是线性时不变，即任何时间的输入相同、输出也相同。为此$(\Delta,A,B,C),\;(\bar{A},\bar{B})$都应该是对任意时间都相同的。

但是为了让S4对输入具有选择性，把$\Delta,B,C$都改造为了输入的线性投影：

$$B=s_B(x)=Linear_N(x),\;C=s_C(x)=Linear_N(x),\\s_{\Delta}=Broadcast_D(Linear_1(x))$$

A因为跟$\Delta$联动所以无需选择性，然后用这个S6模型的一步 + 卷积层 + SwiGLU 结合成一个Mamba层。此外因为无法并行，为了输入长序列的计算速度，设计了一个特殊的卷积核（OpenAI Triton实现）一次性计算所有的$\bar{A},\bar{B}$放在SRAM中、用一次并行累加扫描计算出所有的h、乘上C后才写回HBM。反向传播时则重计算各中间状态以免显存增加。

优点是快，比Attention层计算复杂度仅为线性增长、不会累积Cache节省显存，因此理论上可处理无限序列，并且相比以往的RNN模型记忆力更强，但在长上下文理解、复杂逻辑推理上依旧逊色于最先进的Transformer架构模型，token处理仍无法并行。

Mamba 2在理论上统一了Structured Masked Attention (SMA)和Mamba，并用Tensor Core加速了SSM模块计算，详细暂未了解。

2.5.3.3 一点感想

Mamba将序列模型的基础问题定义为把上下文压缩为一个更小的状态，在此背景下，attention相当于没压缩（因为KV Cache线性增长），传统RNN相当于一直压成一个有限状态。

这很容易让人联想到Agent可能也是一个RNN系统，LLM的上下文窗口就是它的状态budget，用户在离散的无限步骤中每轮都有输入，模型维护一个长期记忆、最近会话历史等等有限状态，这个有限状态的质量越好，输出效果就越好，因此这个RNN的状态转移方法说不定可以用数学方法表示出来。

此外Mamba主要的贡献：Selective of time interval直接决定了RNN状态的更新对新输入（比如token）有多关注，从而可以自适应地记住/忘记输入序列，管理好上下文，这种机制能否迁移到Agent的记忆选取上？

除了记忆的选取机制，结合Sensation and Perception课程的内容，其实遗忘也是人脑处理信息的一个重要机制，它是按什么优先级：发生时间、人、事件、最近被提及的时间、是否是知识……由于目前Agent记忆的增长完全取决于提示词和LLM的能力，记忆的系统机制性遗忘与LLM的处理肯定还是有差别，效果则需要Demo测试。

长期记忆的必要性：或许只体现在它能让Agent更像人：由自己的经历和经验塑造出的多样性极强的智能体，而非每主题刷新。

2.5.4 开源架构拆解

2.5.4.1 Llama 3

2.5.4.2 DeepSeek-V3