最近在给团队做AI芯片的软件栈梳理发现在这个领域待得越久越觉得一颗芯片真正难的不是那颗Die本身而是它外面这圈看不见的软件。编译器、运行时、算子库、驱动、框架适配一层套一层像一座冰山硬件只是露出水面的那个尖。我们决定用Agent帮我们把整条链路的逻辑理成一张“全栈软件地图”从PyTorch里的一个算子调用开始一路追踪到寄存器级怎么写中间经过什么、每层依赖什么、谁在上层等谁、谁在下层被谁等。这篇是系列的第二篇上一篇聊了Agent在芯片项目里能扮演哪些角色这次直接上实操讲我们是怎么用Agent把这张地图画出来的以及画完之后它到底能干什么。1. AI芯片软件栈到底有什么值得拿一张地图去装很多人一听说“全栈”第一反应是前端后端数据库那一套。但是在AI芯片的语境里“全栈”的含义完全不同它指的是从上层AI框架到最底层硬件寄存器的全部软件链路。这张地图如果画得清楚新同学看一周就能定位问题老工程师做性能分析时也能快速判断瓶颈到底落在哪一层。1.1 先把“全栈”这个词拆明白我习惯把AI芯片的软件栈分成七个层次来理解应用层实际的训练脚本、推理服务、模型定义跟具体硬件基本解耦。框架层PyTorch、TensorFlow这些深度学习框架以及相应的前端接口。图编译层负责把模型计算图做算子融合、布局转换、精度选择产出硬件友好的中间表示。中后端编译器针对目标芯片做循环分块、向量化、内存分配、指令调度生成汇编或二进制。运行时层管理设备内存、任务队列、stream并发、同步原语向上承接框架向下对接驱动。驱动与固件处理命令提交、DMA搬移、中断响应、功耗管理这些“脏活累活”。算子库与HAL手写的高性能算子、硬件抽象层把硬件能力封装成相对稳定的API。这七层之间不是简单的线性依赖而是一张网。比如框架层可能绕过图编译直接调用算子库运行时也可能跨过驱动直接操作MMIO寄存器。所以地图不能画成一条直线而是要画成带依赖关系的拓扑结构。1.2 每一层的关键交付物与接口我们做地图时为每一层定义了三个必填的属性做什么事、对外暴露什么接口、依赖下面哪一层的什么能力。这样一张表列下来信息量非常大。我挑几个典型层举例软件层次核心职责关键接口/产物下游依赖框架层承接模型定义与训练循环ATen算子接口、torch.device适配运行时/算子库图编译层图优化、算子选择、内存规划MLIR模块、编译任务文件中后端编译器中后端编译器指令生成、循环变换、寄存器分配可执行ELF/二进制、调试信息驱动加载器运行时层内存管理、stream调度、同步Runtime API、设备上下文驱动/HAL驱动层命令提交、中断、地址映射IOCTL接口、内核模块固件/硬件寄存器这张表看起来简单但把它填准确需要读大量源码和文档。手工填一次至少两周而且填完就过时。这正是不用Agent不行的原因。2. 让Agent来画地图图的是什么我见过不少团队做模块梳理方式是拉一个资深的同事关在小黑屋里写两天文档。听起来很可靠实际上问题很大资深同事的脑子值钱但装不了全仓库的细节新人倒是能翻代码但翻完不理解前因后果。Agent恰好夹在中间——它有耐心读海量文档能跨文件归纳而且永远不会嫌你问的问题蠢。2.1 手工梳理为什么越理越乱先说说我们之前的痛苦。第一个痛点是文档和代码不同步。芯片公司的SDK通常有几十个仓库每一个都有一套自己的README、设计文档、发布说明。仓库A的文档说某接口已经废弃但仓库B的代码还在调驱动代码里的数据结构改了字段名硬件手册却还是老版本。人肉去核对这种差异要花大量时间而且越核越乱。第二个痛点是跨层调用链太长。一个算子的执行路径可能是框架调用算子库 → 算子库申请运行时资源 → 运行时往命令队列里塞一条指令 → 驱动把它翻译成硬件描述符 → 固件轮询执行。这条链路上每一步都有分支和异常处理如果你只盯着其中一层根本看不出整体的性能特征。第三个痛点是新人培养成本高。我见过一个刚入职的同学花了两周才搞清楚“Runtime API”和“Driver API”有什么区别。不是他笨而是没有一张地图告诉他这两个API分别解决什么问题、在什么场景下该看哪个。2.2 Agent能干的和干不了的Agent能干的是那些“读文档、做归纳、照模板输出”的活。它的上下文窗口再长也有限但配合拆解和分批处理它可以把几十万行的代码仓库扫一遍然后按你给定的维度输出结构化的清单。这件事如果让人来做至少要一周而且人的耐心远不如Agent。但Agent干不了判断硬件语义的活。比如“这个同步原语是保证同一队列内有序还是保证不同队列之间的可见性”这种需要结合微架构手册的微妙差异去判断的问题Agent很容易给出看似合理其实是拼凑出来的答案。所以我们的定位是Agent是《聪明的数据整理师》不是最终签字人。地图必须经过人审核才能进库。3. 第一步把芯片的“语料”喂进Agent有人以为用Agent画地图就是把一个问题丢给ChatGPT然后等结果。这么干最多得到一份通用的大道理。真正可用的地图必须基于你自家芯片的一手资料。所以第一步不是问而是喂。3.1 你真正需要的种子文档清单我整理过一份最低限度的文档清单缺了哪一样地图就会有明显空洞ISA手册芯片支持的指令集重点是向量指令、矩阵指令、同步指令。内存模型文档全局内存、共享内存、寄存器堆的层次与一致性语义。TRMTechnical Reference Manual寄存器地址、位域定义、中断号。编译器设计文档图优化pass列表、算子选择规则、循环分块策略。Runtime API文档内存池、stream、event、同步等对象的行为。驱动接口文档IOCTL号、命令描述符格式、IOMMU行为。SDK仓库结构和主要模块说明最好是每个仓库有单独的设计说明。这些文档不需要你提前整理得有多干净。Agent最擅长从一团乱麻里抽丝剥茧但你得保证它能看到原始材料。如果你手里只有一份干巴巴的PPT大纲那再怎么调Prompt也没用。3.2 大文档拆小再用Map-Reduce思路汇总芯片手册动辄上千页Agent的上下文窗口装不下。我的做法是先拆再合思路类似Map-Reduce把手册按主题切成几个分片内存模型、指令集、中断、启动流程。让Agent分别阅读每个分片产出一个几百字的“局部摘要”摘要里必须保留所有接口名和引用编号。把几个局部摘要合并起来再让Agent基于这些摘要做交叉分析找出跨章节的矛盾或依赖。在合并阶段给Agent一个“高维问题清单”例如哪些模块访问共享内存哪些模块负责地址翻译让答案落在地图需要的位置上。这里有个容易犯的错直接把整个仓库塞给Agent让它“总结一下”。结果它会把开头README里的基本信息当成全貌细节全丢。分片阅读虽然费一点token但信息保真度完全不是一个量级。3.3 一个能直接抄的Prompt样例下面是我们跑通的第一版Prompt你可以直接套用。核心是三步设定角色、交代任务、约束输出。你是一位拥有10年经验的AI芯片全栈软件架构师。 你的任务是基于我提供的芯片资料绘制该芯片的软件栈分层地图。 要求 1. 将软件栈划分为7个层次应用框架层、图编译层、中后端编译器、运行时层、驱动层、固件层、硬件抽象层。 2. 对每一层输出以下字段层名称、核心组件、核心职责、对外接口、输入/输出产物、依赖的下一层能力。 3. 列出层与层之间最重要的3条跨层数据流。 4. 对于你不确定的信息使用“待确认”标记不要编造。 5. 用Markdown表格输出并在每个接口名后面标注来源文档ID如[TRM-3.2]。这个Prompt有三个关键点角色要足够资深、输出模板要足够严格、不确定项必须标记。如果没有第三条Agent会用自己的“常识”帮你补全世界观那地图就不是你的芯片地图了而是“所有AI芯片的缝合地图”。4. 第二步让Agent交出一张能当数据库用的分层地图跑完PromptAgent会给你一份看起来很有条理的表格。但这只是初稿距离“能用”还差两步把它转成结构化数据再拉另一双眼睛来审。4.1 输出模板里的每个字段都有用我见过很多团队画架构图画完就贴墙因为图上只有方框和箭头没有可查询的接口信息。我们的地图不同每个组件都带一份“身份证”组件名、所属层、输入、输出、关键接口、依赖项、配置文件路径、环境变量、最常用的调试命令。比如“内存分配器”这个组件身份证大概长这样所属层运行时层对外接口DeviceMemAlloc(size, alignment)、DeviceMemFree(ptr)内部策略首次适配、伙伴系统、内存池复用依赖驱动层的MapDeviceMem()HAL层的mmap实现配置文件runtime/config/memory_config.json调试命令rt_mem --dump-pool这有什么用用处太大了。当线上出现显存泄漏你不需要去翻十几万行代码先查地图里“内存分配器”的调试命令和数据流定位路径能少一半。当你想优化某个算子的显存占用你先看地图里它与哪个分配策略绑定再考虑是不是要换策略。地图的价值不在于好看而在于它把知识沉淀成了可检索的结构。4.2 用YAML把地图变成可程序化消费的数据纯Markdown表格适合人看但程序没法方便地做依赖分析和变更追踪。我们后来把地图转成了YAML每个组件都是一个node每条依赖都是一条edge。Agent对YAML的理解能力很强只要在Prompt里给出示例它就能生成合规的结构。layers: - name: runtime components: - name: memory_pool interfaces: - DeviceMemAlloc - DeviceMemFree dependencies: - driver.map_device_mem - hal.mmap config_path: runtime/config/memory_config.json notes: 支持首次适配与伙伴系统线程安全由互斥锁保证 - name: stream_scheduler interfaces: - StreamCreate - StreamSynchronize dependencies: - runtime.memory_pool - driver.submit_command有了这个YAML我可以写一个脚本自动生成架构图SVG、做拓扑排序、检查哪条依赖断掉了。甚至可以让Agent每周扫描一次代码库自动对比“地图上的接口”和“代码里实际导出的符号”是否一致。不一致的地方就是文档腐化的信号。这一步做完地图就不再是一篇死文档而是变成了一个可以被持续消费的数据资产。4.3 红队审查让另一个Agent挑毛病的价值我逐渐养成了一个习惯初稿生成后绝不直接信。我会用另一个Agent扮演“硬件验证工程师”让它对这份地图做红队审查。审查Prompt大概是你是AI芯片硬件团队的验证负责人。下面是一份软件栈地图初稿。请逐项检查 1. 每个接口是否在硬件手册中有对应依据 2. 数据流路径是否完整是否存在只有发出方没有接收方的连接 3. 有没有混用通用芯片术语比如把“GPU同步原语”直接套到NPU上 4. 哪些描述“看起来合理但缺少硬件依据”请用质疑标注。这招比人审效率高很多因为它不会疲劳。有一回我们做算子下沉流程的地图初稿里写了“算子通过PCIe直接访问系统内存”。红队Agent读了一遍TRM指出该NPU没有PCIe直连能力改成了“通过SMMU映射共享内存”。这个错误如果让人审可能要到集成测试阶段才会暴露。所以我现在坚持“一画一红队”成本不高买的是安全感。5. 从能看到能抄地图工程化之后我踩过的坑地图初版跑通后我们实际用了三个多月期间踩了不少坑。挑几个有代表性的聊聊这些经验比画地图本身更值钱。5.1 术语污染X86的惯性思维会带偏第一个坑是Agent习惯用通用CPU的术语去描述硬件行为。比如让Agent描述“地址翻译流程”它会下意识写“TLB miss之后走页表遍历”。但对于很多AI芯片地址翻译可能依赖专用IOMMU硬件并没有传统意义上的TLB甚至有的芯片只支持连续物理内存根本不需要页表。Agent不知道这些因为它大部分训练语料来自x86世界。解决方法是给Agent喂一份“术语禁忌表”明确告诉它哪些词不可以用换成什么。比如“NPU context”不能说成“CUDA context”“设备内存”不能说成“显存”。这是很傻但很有效的办法能让地图少很多误导。5.2 缺失与幻觉怎么让Agent承认“不知道”第二个坑是幻觉。Agent在信息不足时特别喜欢补全比如地图上某个寄存器定义缺失它会按通用规律给你补一个“看起来合理”的位域。这在地图场景里非常危险因为硬件工程师一旦信了调试时会被带偏。我们后来在Prompt里加了“置信度三级制”高置信度资料中明确写出且多个来源一致。中置信度资料中有相关描述但存在版本差异。低置信度资料中未直接定义由模型推理得出。所有低置信度项必须用待确认标记并且在地图的元数据区域集中列出。这样人审的时候只需要盯着这些待确认项看整个审核时间缩短了大概一半。5.3 持续维护地图需要跟着代码一起长第三个坑是地图会老化。SDK迭代速度很快一个月后接口名可能就变了新加的算子没有进地图废弃的路径没被删掉。如果不做维护地图的价值会随时间衰减成废纸。我们定的维护策略是每周跑一次“地图体检脚本”。脚本会检查两部分一是代码仓库里的符号导出表二是地图YAML里的接口列表。找出新增、删除、变更的符号自动生成差异报告。然后让Agent基于差异报告更新YAML标注变更来源和日期。这个流程跑顺之后地图就变成了活文档团队成员可以放心地长期依赖它。回看这三个月的实践我最深的体感是Agent不是一个替你思考的顾问而是一个执行力极强的整理器它把散落在几十个文档、几万行代码里的信息磨成一份可编辑、可检索、可更新的结构真正吃力的是搭建让它发挥作用的框架。地图画出来只是起点把它接入开发流程、用数据驱动地维护它才是这套打法真正的门槛。如果你正准备梳理自家芯片的软件栈别纠结用什么模型先把手头的TRM和仓库目录理一遍然后让Agent从第一层开始画。