本文浅析了谷歌 TPU 与国内 AI 加速器厂商在 Linux 上的内核驱动与内存HBM管理机制以及它们为何普遍不采用 DRM 框架的原因。目录1. 背景AI 加速器的物理内存2. 谷歌 TPU 的驱动与 HBM 管理3. 国内 AI 加速器厂商驱动栈全景4. 为什么自研厂商普遍不用 DRM5. 走 DRM 路线的例外海光与摩尔线程6. 各路线与各厂商的优劣分析7. 后续有没有可能统一8. 总结1. 背景AI 加速器的物理内存AI 加速器芯片普遍具有多层物理存储结构与 GPU 类似存储层级介质作用典型容量主存板载 HBM片外 DRAM存放模型参数、激活值、中间结果数十 GB 级片上高速缓冲SRAM / VMEM / CMEM暂存矩阵单元MXU计算数据MB 级以谷歌 TPU 为例TPU v4 每芯片约 32GB HBMv5e/v5p 容量更大片内还有向量内存VMEM和矩阵乘法单元缓冲区。因此讨论HBM 管理本质是讨论片外物理 DRAM 的分配、地址映射与主机-设备数据搬运由哪一层软件负责。2. 谷歌 TPU 的驱动与 HBM 管理2.1 Cloud TPU数据中心v2/v3/v4/v5内核驱动私有gasketapex框架非标准 DRM。GasketGoogle ASIC Software, Kernel Extensions and Tools通用内核框架为 PCIe 挂载的 ASIC 提供字符设备、DMA、中断、BAR 映射等基础设施曾出现在drivers/staging/gasket/。Apex构建在 Gasket 之上的具体 TPU 设备驱动。HBM 管理主要在用户态由libtpu配合 XLA / TensorFlow / JAX runtime完成。内核驱动只负责暴露设备、DMA 映射、命令队列提交。HBM 地址空间划分、buffer 分配、调度由XLA 编译器 runtime决定通过 DMA 搬入/搬出 HBM。2.2 边缘 TPUCoral / Edge TPU使用同一套Gasket Apexgasket.koapex.ko。用户态通过libedgetpu访问。2.3 与 GPU 内存管理路径的对比维度AMD / NVIDIA GPU谷歌 TPU内核子系统DRM / TTM / GEMGasket非 DRM显存 / HBM 管理内核 TTM buffer manager主要在用户态libtpu / XLA内存迁移DRM / HMM /migrate_vma由 runtime 显式 DMA简述Cloud TPU 的 HBM 通过谷歌私有的gasket/apex内核驱动暴露设备并做 DMA而 HBM 的实际分配与管理主要由用户态的 libtpu/XLA runtime 负责不走 Linux 标准的 DRM/TTM 显存管理路径。3. 国内 AI 加速器厂商驱动栈全景国内厂商绝大多数走自研私有内核驱动 私有用户态 runtime风格接近 NVIDIA CUDA 闭源栈或谷歌 Gasket 模式。3.1 华为昇腾AscendNPU / 达芬奇架构内核驱动私有模块如davinci、devmmdevice memory management、hdc、dpc设备节点/dev/davinciX、/dev/devmm_svm。HBM 管理设备内存由内核驱动 用户态 runtime 管理支持SVM统一虚拟内存devmm负责 host/device 地址空间映射思路类似 HMM但自研。用户态CANN 栈ACL / Runtime / GE / AICPU分配接口如aclrtMalloc。是否 DRM否。3.2 寒武纪CambriconMLU内核驱动私有cambricon_drv.ko设备节点/dev/cambricon_devX、/dev/cambricon_ipcm。内存管理自研设备内存分配器用户态通过 CNRT / CNDrvcnrtMalloc分配支持 host-device 统一地址。用户态NeuwareCNToolkit / CNRT / CNNL / CNCL。是否 DRM否。3.3 燧原科技 EnflameGCU最接近TPU定位内核驱动私有enflame/gcu内核模块。用户态TopsRider 栈TopsRuntime对接 TensorFlow / PyTorch / XLA。特点采用 XLA 后端软件路径与谷歌 TPU 最相似HBM 管理在用户态 runtime 内核 DMA。3.4 百度昆仑芯Kunlun / XPU内核驱动私有kunlun/xpu。用户态XREKunlun Runtime Environment、XDNN接口如xpu_malloc。是否 DRM否。3.5 壁仞 Biren / 天数智芯 Iluvatar / 沐曦 MetaX / 摩尔线程GPGPU 类“类 CUDA”壁仞私有内核驱动 BIRENSUPA 软件栈对标 CUDA。天数智芯私有驱动 类 CUDA 的 Corex / IXUCA。沐曦 MetaX私有内核驱动metax/maca模块 MACA / MXMACA软件栈主打CUDA 兼容与迁移编译器mxcc算子库 mcBLAS/mcDNN/mcFFT 对应 cuBLAS/cuDNN/cuFFT拥有曦云 MXC训练/通用计算、曦思 MXN推理、曦彩 MXG图形三条产品线。计算主通道为私有栈类 NVIDIAnvidia.ko模式。摩尔线程mtgpu驱动基于 DRM 框架需同时做图形 计算显存走 TTM/GEM用户态 MUSA 对标 CUDA。3.6 海光 DCUHygon特殊源自 AMD 授权软件栈基于 ROCm 分支DTK内核驱动是amdgpu/amdkfd 的定制版本走DRM KFD路径。与本仓库rocr-runtime / amdgpu SVM技术栈直接同源。3.7 汇总对比表厂商内核驱动是否走 DRM用户态栈HBM / 显存管理华为昇腾私有 davinci/devmm否CANN自研 SVM寒武纪cambricon_drv否Neuware / CNRT自研分配器燧原enflame / gcu否TopsRider (XLA)runtime DMA昆仑芯kunlun / xpu否XRE / XDNN自研壁仞私有否BIRENSUPA自研海光 DCUamdgpu/kfd 定制是DRMKFDDTK / ROCmTTM HMM/SVM核心规律纯 AI 加速器昇腾、寒武纪、燧原、昆仑→ 私有字符设备驱动 私有 runtimeHBM 管理在内核私有模块 用户态不用 DRM/TTM。图形 GPU 或源自 AMD 的摩尔线程、海光→ 走 DRM其中海光与 amdgpu SVM / rocr-runtime 是同一套 KFD HMM 内存迁移机制。4. 为什么自研厂商普遍不用 DRM本质上是技术需求 工程成本 商业策略的三方权衡。4.1 DRM 是为图形显示设计的AI 加速器不需要DRMDirect Rendering Manager的抽象体系围绕 GPU 图形渲染建立KMS显示模式设置、framebuffer、CRTC、连接器、EDID、vblank、显示 fence 同步……GEM/TTM 中大量概念scanout buffer、tiling、显示扫描对纯计算芯片无意义。纯 AI NPU/TPU没有显示输出用不上 DRM 约 90% 的功能。为了剩余 10% 的 buffer 管理背负整个 DRM 框架不划算。4.2 DRM/TTM 太重、太复杂TTM 的内存迁移、驱逐、多域VRAM/GTT/system管理是为显存有限、要和 CPU 抢内存、支持图形负载设计的逻辑极其复杂。AI 芯片内存模型更简单HBM 是一大块线性地址空间runtime 自己做 bump / buddy 分配器即可不需要 TTM 的驱逐搬迁机制。一个数百行的字符设备驱动ioctl mmap DMA即可满足需求远比接入 TTM 简单。4.3 上游 DRM 有严格的社区规则和维护负担必须开源用户态栈DRM maintainer 明确不接受只有闭源用户态的 DRM 驱动。代码风格、UAPI 稳定性、review 流程严格周期以年计。一旦进主线UAPI 需永久向后兼容束缚硬件快速迭代。对追求快速出货、且不愿开源核心 runtime/编译器的厂商这不可接受故选择out-of-tree 私有字符设备驱动。4.4 商业机密与生态锁定AI 芯片核心竞争力在编译器 runtime 内存调度策略走 DRM 意味着暴露内存管理 UAPI 细节。私有栈可将 driver / runtime / compiler 打包为闭源 SDKCANN、Neuware、TopsRider…形成对标 CUDA 的生态壁垒。4.5 参考对象是 CUDA不是 MesaNVIDIA 计算栈本身不走 DRMnvidia.ko是私有字符设备驱动只有nouveau才是 DRM。国内厂商照抄这套成熟商业模式私有内核模块 /dev/xxx字符设备 闭源 runtime。5. 走 DRM 路线的例外海光与摩尔线程厂商采用 DRM 的原因海光 DCU源自 AMD 授权直接继承amdgpu KFD代码改比重写省事摩尔线程做真正的图形 GPU需要显示输出 / 渲染DRM/KMS 是刚需这两家要么被迫继承要么确实需要图形才走 DRM 路线。6. 各路线与各厂商的优劣分析6.1 两条路线的优劣路线 A私有字符设备驱动 闭源 runtime代表昇腾 / 寒武纪 / 燧原 / 昆仑 / NVIDIA / 谷歌 TPU优势迭代快UAPI 自定硬件换代不受上游兼容性约束改内存模型/指令集不用管社区。实现轻不背 DRM/TTM 的图形包袱几百行字符设备ioctl mmap DMA即可跑通。护城河driver runtime compiler 打包成闭源 SDK形成 CUDA 式生态锁定保护核心 IP。内存模型简单可控HBM 当线性空间自己做分配器调度策略完全自定义。劣势生态碎片化每家一套 APICANN / Neuware / TopsRider…互不兼容用户迁移成本高。无法进主线永远是 out-of-tree 模块随内核升级易 break需厂商持续适配。黑盒难调试闭源出问题依赖厂商社区帮不上忙。重复造轮子内存管理、DMA、IOMMU 等基础设施每家各写一遍质量参差。安全审计难闭源内核模块是攻击面云厂商/客户难以信任。路线 BDRM TTM/GEM代表摩尔线程 / 海光 / AMD / Intel优势复用成熟基础设施TTM 内存驱逐、GEM buffer 共享、dma-buf、fence 同步、IOMMU 集成全都现成。能进上游驱动进 mainline 后随内核长期维护发行版开箱即用。图形 计算统一一套栈同时支持渲染和 compute摩尔线程刚需。标准互操作dma-buf / PRIME 让跨设备零拷贝共享、与显示子系统协作变简单。劣势必须开源用户态社区硬性要求核心 runtime/compiler 难以闭源商业机密受限。框架重、门槛高TTM 复杂度高接入和 debug 成本大对纯 AI 芯片是过度设计。UAPI 永久兼容进主线后接口要长期向后兼容束缚硬件激进创新。review 周期长合入以年计不利于快速出货。6.2 具体厂商的优劣厂商优势劣势谷歌 TPUXLA 编译器成熟、软硬协同极致、规模化部署完全私有、只能在谷歌云用、无对外生态华为昇腾国产最完整栈(CANN)、支持 SVM、生态投入大API 学习曲线陡、闭源、迁移成本高寒武纪起步早、Neuware 相对完整生态小、市占低、闭源黑盒燧原走 XLA/OpenXLA路径通用、易接标准框架体量小、软件成熟度待验证昆仑芯背靠百度内部大规模场景验证对外生态弱、XPU 编程模型小众壁仞 / 天数类 CUDA、迁移门槛相对低供应链风险、驱动闭源沐曦 MetaXCUDA 兼容激进、迁移门槛低、训练推理图形全产品线驱动/runtime 闭源、生态追赶中、供应链风险摩尔线程唯一 DRM 图形 计算通吃、MUSA 对标 CUDA计算性能与生态仍追赶中、DRM 包袱海光 DCU直接复用 ROCm/amdgpu生态最省力、SVM/HMM 现成依赖 AMD 授权、架构受制于 AMD 迭代节奏6.3 小结纯 AI 芯片选私有栈是用生态封闭换迭代自由与商业护城河选 DRM是用必须开源、框架沉重换上游维护与标准互操作。对国产而言昇腾代表自成体系的封闭强栈路线海光代表复用 AMD 开源栈的省力路线——后者恰好与本仓库的amdgpu SVM / rocr-runtime同源SVM/HMM 内存迁移几乎零成本继承。7. 后续有没有可能统一统一的可能性是有的但会分层次、分阵营地进行而不会全球收敛到单一栈。7.1 内核驱动层正缓慢向统一基础设施靠拢Linux 社区正在把计算加速器从 DRM 里抽象出来drivers/accel/子系统Accelerator subsystem2022 年进主线专为不做图形的 AI/计算加速器设立复用部分 DRM 基础设施GEM、drm_device、dma-buf、fence但剥离图形/KMS 部分。已有 HabanaIntel Gaudi、Intel VPUivpu、AMDamdxdnaRyzen AI NPU等接入。这恰好解决了第 4 章的DRM 太重、图形包袱痛点accel 给纯 AI 芯片一个轻量版 DRM的上游归宿。潜在统一点未来国产厂商若想进主线、被发行版开箱支持drivers/accel/是最现实的路径但目前国产厂商基本仍在 out-of-tree动力不足。7.2 DRM 自身的模块化drm_gpuvm/drm_gpusvm/drm_pagemap需要特别指出“DRM 图形专用重框架”的印象已经过时。近两年 DRM 正在长出一批跨驱动共享的通用内存/地址空间管理中间层把过去各家私有的逻辑收敛成公共组件drm_gpuvmGPU 虚拟地址空间管理器前身为drm_gpuva_mgr2023 年进主线。把“GPU VA 空间 BO 映射区间drm_gpuva”的通用管理逻辑区间树、split/merge、VM_BIND 语义抽出。采用者Nouveau首个、Xe、Panthor、PowerVR 等正成为新驱动的标配。drm_gpusvmGPU 共享虚拟内存2024 年随Xe引入目标就是做跨驱动可复用的 SVM 基础设施。基于HMM的 system allocator——CPU 与 GPU 共享同一虚拟地址空间。drm_pagemap配合drm_gpusvm管理 device-private 的ZONE_DEVICE内存、用migrate_vma做页面迁移把过去 amdgpusvm_range、Nouveau 等各自实现的 SVM 迁移逻辑收敛成公共层。这意味着 DRM 正从“图形专用重框架”演化为“模块化、可按需取用的 GPU 基础设施库”能力老做法各驱动私有新的 DRM 公共层GPU 地址空间 / VM_BIND各写 VA 管理drm_gpuvmSVM / 统一内存迁移amdgpusvm_range、nouveau 各写drm_gpusvmdrm_pagemapBO 内存管理—TTM已有对两条路线的影响走 DRM 的厂商AMD、Intel Xe、海光能直接复用这套 SVM 基础设施而走私有栈的纯 AI 厂商昇腾等依然在自研 SVM享受不到这层红利——这反而拉大了两条路线在“统一虚拟内存能力”上的差距。7.3 用户态编程层更可能通过编译器 IR 框架后端实现事实统一内核难统一但用户态正在被上层抹平统一层机制现状框架后端PyTorch 的PrivateUse1/torch.compile、OpenXLA/StableHLO各家写后端即可接入用户不感知底层编译器 IRMLIR / OpenXLA / Triton燧原、部分国产已走 XLA/MLIR中间接口SYCL / oneAPI、OpenAI Triton试图做跨厂商 CUDA 替代趋势绝大多数厂商都在做PyTorch 后端 XLA/MLIR 接入。用户写 PyTorch不再关心是昇腾还是海光——这就是事实统一但底层驱动/runtime 仍各自私有。7.4 为什么完全统一很难商业护城河CUDA 的成功恰恰在于不统一锁定。厂商没有动力交出 runtime/编译器控制权。地缘/供应链国产阵营昇腾、海光…与 NVIDIA/CUDA 阵营被动脱钩反而会形成两套甚至多套并行标准。硬件架构差异大NPU脉动阵列、GPGPUSIMT、达芬奇Cube指令模型差异根本底层 UAPI 难以真正统一。UAPI 永久兼容成本进主线后接口冻结厂商顾虑迭代自由第 4.3 节提过。7.5 最可能的结局分层收敛而非单点统一用户态框架层 → 高度统一PyTorch / OpenXLA 抹平差异 ★最可能 ↑ 编译器 IR 层 → 部分统一MLIR / StableHLO 成公约数 ↑ Runtime/驱动层 → 阵营内可能统一全球难统一 ↑ 内核子系统 → drivers/accel drm_gpuvm/drm_gpusvm/drm_pagemap 提供可选统一底座走 DRM 的厂商可直接复用7.6 小结统一在上层PyTorch/OpenXLA 框架后端和内核层DRM 公共基础设施同时推进底层 runtime 则分阵营并存。一方面 PyTorch/OpenXLA 在上层抹平差异另一方面内核层drivers/accel/与drm_gpuvm/drm_gpusvm/drm_pagemap正把 VM 与 SVM 能力做成跨驱动公共组件。但由于商业锁定、地缘脱钩和架构差异最现实的未来仍是上层框架 内核基础设施事实统一、厂商 runtime 分阵营并存——类似今天 CPU 世界有统一的 C/POSIX但各家微架构各不相同。对国产阵营而言海光走 ROCm/DRM 已经天然贴近上游统一底座能直接受益于drm_gpusvm/drm_pagemap等新基础设施而昇腾等封闭栈既享受不到内核 SVM 公共层又更依赖PyTorch 后端这条上层统一路径。8. 总结DRM 是为图形 GPU 设计的重型框架。纯 AI 加速器既不需要显示功能又想保护闭源 runtime 并快速迭代因此选择轻量的私有字符设备驱动 闭源用户态栈——这条路更简单、更自由、也更符合对标 CUDA 的商业策略。只有需要图形输出或继承自 AMD/Intel 现有 DRM 代码的厂商才会选择 DRM。谷歌 TPU 用私有gasket/apexHBM 管理放在用户态 libtpu/XLA。国内纯 AI 加速器昇腾、寒武纪、燧原、昆仑均为私有驱动 私有 runtime。海光 DCU 是唯一与本仓库 amdgpu SVM / rocr-runtime 技术同源的方案DRM KFD HMM/SVM 内存迁移。注各厂商驱动源码大部分未开源模块名与实现细节以厂商官方发布为准本文基于公开资料整理。