ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

AI芯片全栈软件地图:用Agent梳理UMD、KMD与运行时

AI芯片全栈软件地图:用Agent梳理UMD、KMD与运行时 1. 从一块芯片说起为什么全栈软件地图比芯片本身更难画我入行做芯片软件栈那几年最大的感受就是芯片设计出来只是万里长征第一步真正让一块 AI 芯片跑起来、跑得稳、跑得高效靠的是背后一整套软件栈的协同。很多人以为做 AI 芯片就是堆算力、拼制程、比 TOPS 数字但实际上一块芯片从流片回来到客户能跑通第一个模型中间隔着的软件工作量往往比硬件设计本身还要大。这篇文章要聊的就是怎么用 Agent 的思路把一块 AI 芯片的全栈软件地图给梳理清楚。所谓“全栈软件地图”说白了就是从最底层的驱动、内核态模块到中间的运行时、编译器、算子库再到上层的框架适配、模型部署工具链这一整条链路上到底有哪些组件、它们之间怎么协作、每一层的核心职责是什么。而“用 Agent 做”这件事指的是借助 AI Agent 的能力去自动化地理解、梳理、甚至生成这套软件栈的文档、依赖关系和调试路径。为什么这件事值得单独拿出来讲因为 AI 芯片的软件栈有一个非常典型的特点层次多、耦合深、文档散。一个典型的 AI 芯片软件栈从下往上至少包括 UMD用户态驱动、KMD内核态驱动、运行时库、编译器后端、算子库、框架插件这几大块每一块又可能对应好几个代码仓库、好几套构建系统、好几种调试工具。新人进来想搞清楚“一个算子从框架调用到最终在芯片上执行中间经过了哪些层”没有一张清晰的地图基本就是抓瞎。这张地图适合谁来参考如果你是刚进入 AI 芯片公司的软件工程师它能帮你快速建立全局认知如果你是做 Agent 开发的它能给你一个真实的、复杂的软件系统梳理场景如果你是技术管理者它能帮你理解软件栈的复杂度到底体现在哪里从而更合理地排期和分工。接下来我会从整体设计思路、核心组件拆解、实操梳理流程、常见坑排查这几个维度把这张地图一层层铺开。2. 全栈软件地图的整体设计与分层思路2.1 为什么 AI 芯片软件栈要分层而不是一坨先讲一个我踩过的坑。早年参与一个项目团队为了赶进度把驱动、运行时、算子库的接口混在一起写结果就是任何一个模块改动都要重新编译整个栈调试的时候根本分不清问题出在哪一层。后来痛定思痛重新做分层才把维护成本降下来。AI 芯片软件栈分层的核心逻辑和操作系统分层是一个道理每一层只暴露必要的接口给上一层隐藏下层的实现细节。这样做的好处有三个。第一是解耦上层框架不需要知道底层芯片的寄存器怎么配只需要调用运行时提供的统一接口。第二是可替换换一代芯片只要运行时接口不变上层框架和算子库基本不用动。第三是可调试出了问题可以按层定位是驱动没起来还是运行时调度错了还是算子实现有 bug。一个典型的 AI 芯片全栈软件分层从下到上大致是这样的层级组件核心职责典型产物内核态KMD设备管理、内存映射、中断处理内核模块 .ko用户态驱动UMD命令队列、任务提交、设备控制动态库 .so运行时Runtime内存管理、流/队列调度、事件同步运行时库编译器Compiler图优化、算子融合、代码生成编译产物算子库Kernel Library高性能算子实现算子二进制框架适配Framework Plugin对接 PyTorch/TensorFlow 等框架插件这张表看起来简单但真正落地的时候每一层的边界在哪里、接口怎么定是要反复打磨的。我个人的经验是UMD 和 KMD 的边界要尽量清晰KMD 只做最必要的硬件操作和资源管理把策略性的东西都放到 UMD这样内核态代码量小、稳定、好维护。而运行时和编译器的边界则要看你的芯片是偏通用还是偏专用专用芯片可以把更多优化放到编译器里。2.2 用 Agent 来梳理地图到底解决了什么问题传统梳理软件栈地图的方式无非是读代码、看文档、问老人。但这三个方式在 AI 芯片场景下都有明显的短板。读代码代码量动辄几十万行跨多个仓库人脑很难建立全局关联。看文档文档往往滞后于代码而且散落在各个 wiki 里找起来费劲。问老人老人时间有限而且很多隐性知识他自己都说不清楚。Agent 在这里的价值是它能自动化地做信息聚合和关系抽取。具体来说你可以让 Agent 去扫描指定的一组代码仓库提取出每个模块的对外接口、依赖关系、构建脚本然后自动生成一张模块依赖图。你还可以让 Agent 去读散落的文档把关于同一个组件的描述聚合到一起形成一份统一的说明。更进一步你可以让 Agent 去分析调用链回答“这个算子从框架到硬件经过了哪些函数”这类问题。但这里要泼一盆冷水Agent 不是万能的它做的是辅助梳理不是替代理解。Agent 生成的依赖图可能有遗漏聚合的文档可能有冲突调用链分析可能因为动态派发而断掉。所以正确的用法是把 Agent 当成一个不知疲倦的初级助手它帮你把 80% 的机械性信息收集工作做完剩下 20% 的判断和验证还是得靠人。我在实际项目里通常会让 Agent 先跑一遍生成初稿然后我拿着初稿去和代码对照把错的改掉、漏的补上这样效率比纯手工高很多。2.3 地图的粒度怎么定太粗没用太细爆炸定粒度是个技术活。我见过有人把地图做到函数级别结果一张图几千个节点根本没法看。也见过有人只画到模块级别结果真出问题的时候还是不知道该看哪个文件。我的经验是地图的粒度要跟着使用场景走。如果是给新人做入门导览画到模块级别就够了重点是让他知道有哪几大块、每块大概干什么。如果是给调试用那就要做到关键接口级别比如 UMD 的任务提交接口、运行时的内存分配接口这些是问题高发区。如果是给性能优化用那可能要深入到算子级别看每个算子的实现和调度策略。用 Agent 做的时候可以分层次生成。先让 Agent 生成模块级的地图作为骨架。然后针对重点模块再让 Agent 深入到接口级甚至函数级生成细节地图。这样既有全局又有局部不会一上来就被细节淹没。我一般会维护三份地图一份模块级的总览图一份关键路径的接口图一份高频问题的排查图三份配合使用。3. 核心组件拆解UMD、KMD 与运行时的关键细节3.1 UMD 用户态驱动任务提交的咽喉要道UMD 是整个软件栈里我最关注的一层因为它是用户态和内核态的交界也是任务提交的咽喉。一个 AI 芯片的 UMD核心职责包括接收上层运行时的任务描述把它翻译成硬件能理解的命令通过 KMD 提供的接口提交给硬件然后管理任务的完成状态。这里的关键设计点是命令队列的管理。AI 芯片通常有多个计算单元任务需要被分发到不同的队列上并行执行。UMD 要负责队列的创建、任务的入队、优先级的调度、以及队列之间的同步。我踩过的一个坑是早期队列数量设得太少导致高优先级任务被低优先级任务堵住推理延迟抖动很大。后来改成多级队列加优先级抢占才把尾延迟压下来。用 Agent 梳理 UMD 的时候重点让它提取这几类信息UMD 对外暴露的 API 列表、每个 API 的参数和返回值、API 内部的调用链、以及 UMD 和 KMD 之间的 ioctl 接口。这些信息合起来就能画出一张 UMD 的功能地图。我实测下来Agent 对 API 列表和调用链的提取准确率比较高但对 ioctl 这种涉及结构体定义的接口容易漏字段需要人工补。提示梳理 UMD 时一定要把错误码体系单独拎出来。AI 芯片的 UMD 错误码往往有几十个每个对应不同的失败场景这是排查问题的关键线索。3.2 KMD 内核态驱动稳定压倒一切KMD 是跑在内核态的它的第一原则是稳定第二原则还是稳定。因为内核态代码一旦出问题轻则进程崩溃重则整机挂死。所以 KMD 的设计哲学是尽量薄、尽量简单、尽量少做决策。KMD 的核心职责包括设备初始化、内存的物理映射、中断处理、以及为 UMD 提供 ioctl 接口。其中内存映射是重头戏AI 芯片的显存或者叫设备内存管理涉及到物理页分配、页表建立、缓存一致性等一堆细节。我见过因为缓存一致性没处理好导致算出来的结果偶尔错一位的 bug查了整整两周。用 Agent 梳理 KMD 的时候要特别注意它的并发模型。KMD 里的代码可能被多个进程、多个线程同时调用哪些数据结构有锁保护、哪些是无锁的、中断上下文和进程上下文怎么同步这些是理解 KMD 的关键。Agent 可以帮你把锁的使用点列出来但锁的正确性判断还是得靠人。我一般会让 Agent 生成一份“锁清单”然后逐个 review。3.3 运行时承上启下的调度中枢运行时是软件栈里最“忙”的一层它上接框架下接驱动中间还要管内存、管流、管事件。一个设计良好的运行时应该让上层感觉不到底层硬件的复杂性同时又能充分发挥硬件的性能。运行时的核心组件包括内存管理器、流/队列调度器、事件同步机制、以及算子分发器。内存管理器负责设备内存的分配和回收这里有个经典问题是内存池的设计池子太小会频繁分配释放池子太大又浪费显存。我的经验是内存池要按块大小分级小块和大块分开管理减少碎片。流调度器负责把任务分配到不同的流上并行执行流之间的依赖通过事件来同步。这里的关键是依赖关系的正确表达如果依赖漏了会出现数据竞争如果依赖多了会损失并行度。用 Agent 梳理运行时的时候可以重点让它分析流和事件的创建、使用、销毁路径把同步关系画出来。这张图对理解性能瓶颈特别有用。3.4 编译器与算子库性能的主战场编译器和算子库是决定芯片实际性能的地方。编译器负责把框架的图转化成芯片能执行的指令序列算子库则提供了高性能的算子实现。这两块的技术深度很大我这里只讲和地图梳理相关的部分。编译器这块重点是理解它的编译流程和优化 pass。一个典型的 AI 编译器会经过图解析、算子融合、内存规划、指令生成这几个阶段。每个阶段都有对应的 passpass 之间通过中间表示IR传递。用 Agent 梳理的时候可以让它把 pass 列表和每个 pass 的作用提取出来形成一张编译流程图。算子库这块重点是理解算子的覆盖范围和实现策略。哪些算子有手写的高性能实现哪些是走通用模板生成的哪些是 fallback 到 CPU 的这些信息对性能调优很重要。Agent 可以帮你把算子列表和对应的实现文件关联起来生成一张算子覆盖表。4. 实操过程用 Agent 一步步生成软件地图4.1 第一步确定扫描范围和输入源动手之前先要把输入源理清楚。AI 芯片软件栈的代码通常分散在多个仓库里比如驱动一个仓库、运行时一个仓库、算子库一个仓库、框架插件一个仓库。你要先确定这次梳理覆盖哪些仓库把它们 clone 到本地或者让 Agent 能访问到。除了代码还要把文档、构建脚本、CI 配置也纳入输入源。文档能提供设计意图构建脚本能揭示模块依赖CI 配置能反映测试覆盖。我一般会准备一个清单把所有输入源列出来然后让 Agent 逐个处理。这里有个实操技巧给 Agent 的输入要分优先级。核心模块的代码是必读的边缘模块的代码可以只读接口。文档里设计文档优先于使用文档。这样能保证 Agent 在有限的处理能力下优先把最重要的信息提取出来。4.2 第二步让 Agent 提取模块和接口信息输入源准备好之后就可以让 Agent 开始提取信息了。这一步的目标是生成一份结构化的模块清单每个模块包含模块名、所在仓库、对外接口、依赖的其他模块。我通常会给 Agent 这样的指令扫描指定目录下的头文件和源文件提取所有对外暴露的函数和类分析它们之间的调用关系输出一份 JSON 格式的模块描述。Agent 处理完之后我会抽查几个模块看看提取的接口全不全、依赖关系对不对。实测下来Agent 对 C/C 代码的接口提取效果比较好对 Python 代码的装饰器、动态属性这些特性容易漏。所以如果你的软件栈里有 Python 层要额外提醒 Agent 注意这些地方。另外宏定义和模板元编程也是重灾区Agent 经常看不懂需要人工补。4.3 第三步生成依赖图和调用链模块清单有了之后下一步是生成依赖图和调用链。依赖图展示模块之间的静态依赖关系调用链展示一次任务从上层到下层的动态执行路径。依赖图可以让 Agent 根据模块清单自动生成用 Graphviz 或者类似的工具渲染出来。调用链则需要 Agent 去分析具体的代码路径比如从框架的算子调用开始一路追到 UMD 的任务提交。这一步的难点在于动态派发和回调很多地方是通过函数指针或者虚函数调用的静态分析追不下去。我的做法是让 Agent 把能追的追出来追不下去的地方标记出来然后人工补上。调用链的价值在于它能帮你快速定位问题。比如推理结果不对你可以沿着调用链一层层往下查看是哪一层的输出和预期不符。我一般会把关键路径的调用链保存下来作为排查问题的参考。4.4 第四步人工校验和补充Agent 生成的地图一定要经过人工校验。校验的重点有三个一是接口是否完整二是依赖是否正确三是调用链是否有断点。校验的方法我一般是用几个典型场景去验证。比如选一个常见的算子手动走一遍从框架到硬件的路径看看地图上画的路径和实际是否一致。再比如选一个已知的 bug看看地图能不能帮你定位到问题模块。如果地图在这些场景下都能用那基本就靠谱了。校验过程中发现的错误和遗漏要及时反馈给 Agent让它修正。如果 Agent 支持增量更新那就把修正后的信息喂回去让它重新生成。如果不支持那就人工改地图。这个过程可能要迭代几轮但每轮都会让地图更准。4.5 第五步地图的维护和更新地图不是一次性的产物代码在变地图也要跟着变。所以最后一步是建立地图的维护机制。我的做法是把地图生成脚本化每次代码有大的变更就重新跑一遍 Agent生成新的地图然后和旧地图做 diff看看哪些地方变了。这样既能保证地图的时效性又能通过 diff 发现一些意外的变更。另外地图的存储也很重要。我一般会把地图放在团队的 wiki 上配上搜索功能方便大家查询。地图的每个节点都链接到对应的代码文件点进去就能看源码。这样地图就不只是一张图而是一个可导航的知识库。5. 常见问题与排查技巧实录5.1 Agent 提取信息不全怎么办这是最常见的问题。Agent 提取信息不全通常有几个原因一是输入源没给全二是 Agent 的处理能力有限三是代码里有它不认识的语法。排查思路是这样的先检查输入源确认该给的代码和文档都给了。然后看 Agent 的日志看它处理到哪一步卡住了是不是有文件解析失败。如果是语法问题那就把相关代码片段单独拎出来用更明确的指令让 Agent 处理。我踩过的一个坑是Agent 处理大文件时会截断导致后半部分的接口没提取到。后来我把大文件拆成小文件或者分段喂给 Agent问题就解决了。所以控制单次输入的大小是个实用的技巧。5.2 依赖关系画错了怎么定位依赖关系画错往往是因为 Agent 把间接依赖当成了直接依赖或者把条件编译的依赖也算进去了。定位的方法是挑几条可疑的依赖边手动去代码里验证。比如地图上说 A 依赖 B你就去 A 的代码里搜 B 的头文件看是不是真的 include 了。如果没 include那这条依赖就是错的。条件编译的依赖要看编译时的宏定义这个 Agent 经常搞错需要人工根据构建配置来判断。5.3 调用链断在动态派发处怎么办动态派发是调用链分析的老大难。虚函数、函数指针、回调这些都会让静态分析断掉。我的处理方式是对每个断点人工去分析可能的实现。比如一个虚函数调用你就去找所有 override 这个虚函数的子类把它们的实现都列出来作为可能的分支。函数指针的话就去找所有赋值给这个指针的地方。这样虽然不能做到 100% 自动但能把断点补上。5.4 地图更新后和旧版对不上怎么办地图更新后对不上通常是因为代码重构了模块拆分或者合并了。这时候不要慌先看 diff搞清楚哪些模块变了。如果是拆分就把旧模块的接口分到新模块上如果是合并就把旧模块的接口合到一起。我一般会保留地图的版本历史每次更新都记一笔说明改了什么、为什么改。这样回溯的时候有据可查。5.5 常见问题速查表问题现象可能原因排查方法解决技巧接口提取不全输入源缺失或文件过大检查输入源查看 Agent 日志拆分大文件分段处理依赖关系错误间接依赖误判或条件编译手动验证可疑依赖边结合构建配置判断调用链断裂动态派发找虚函数子类或指针赋值点人工补全分支地图版本混乱代码重构对比新旧地图 diff保留版本历史性能瓶颈难定位地图粒度不够深入到算子级生成算子级细节地图提示排查问题时优先看 UMD 和运行时的交界处这里是问题高发区。很多看似是算子的问题实际是任务提交或内存管理出了问题。6. 我在这件事上的一些个人体会做 AI 芯片全栈软件地图这件事我最大的体会是地图的价值不在于画得多漂亮而在于用起来顺不顺手。我见过很多团队花大力气画了一张精美的架构图结果没人看因为图上的信息和实际排查问题对不上。真正有用的地图是那种你遇到问题能立刻查、查了能立刻定位到代码的地图。用 Agent 来做这件事确实能省很多力气但前提是你要会用它。我的经验是把 Agent 当成一个需要明确指令的助手指令越具体输出越靠谱。不要指望它一次就给你一张完美的地图而是要迭代先出粗稿再逐步细化。另外地图的维护比生成更重要。代码天天在变地图如果不同步更新很快就会过时过时的地图比没有地图更危险因为它会误导人。所以一定要把地图更新纳入日常流程比如每次发版前跑一遍生成脚本确保地图和代码一致。最后分享一个小技巧在地图的每个关键节点上标注一个“常见问题”的链接指向该模块历史上出过的问题和解决方法。这样地图就不只是导航还是一个经验库。这个习惯我坚持了好几年团队里新人的上手速度明显快了很多。
返回列表