VibeTensor: System Software for Deep Learning, Fully Generated by AI Agents

作者/机构: Bing Xu, Terry Chen, Fengzhe Zhou, Tianqi Chen, Yangqing Jia, Vinod Grover, Haicheng Wu, Wei Liu, Craig Wittenbrink, Wen-mei Hwu, Roger Bringmann, Ming-Yu Liu, Luis Ceze, Michael Lightstone, Humphrey Shi NVIDIA

A1 主要贡献

本文探讨了一个核心系统问题:编码智能体(coding agents)是否能够生成一个连贯的深度学习系统软件栈,该软件栈跨越从Python和JavaScript API到C++运行时组件和CUDA内存管理的抽象边界,并仅通过构建和测试进行验证。为了实证回答这一问题,本文开源了VIBETENSOR,这是一个由大型语言模型(LLM)驱动的编码智能体在高级人类指导下完全生成的开源研究性深度学习系统软件栈。“完全生成”指代码来源:所有实现变更均由智能体以差异(diff)形式提出并应用;验证则依赖于智能体工作流执行的构建、测试和差异检查,无需人工逐一审查代码变更。本文的目标并非推出新的训练框架或宣称达到顶尖的端到-端性能,而是聚焦于生成系统的架构、生产和验证方法,以及从中得到的经验教训。

平台支持:该版本主要面向Linux x86_64平台和通过CUDA支持的NVIDIA GPU;目前有意禁用了无CUDA的构建。本文报告了在NVIDIA H100(Hopper, SM90)和Blackwell级GPU上的运行结果。多GPU实验结果仅限于Blackwell平台,并依赖于一个可选的基于CUTLASS的环形allreduce插件,该插件需要CUDA 13+和sm103a工具链支持。

核心贡献
* 端到端系统软件:本文描述了一个深度学习运行时,它拥有多语言前端(Python和Node.js),并由一个共享的C++核心提供支持,该核心实现了调度器、张量/存储实现、自动求导引擎和CUDA子系统。
* 互操作性与扩展点:文档化了DLPack【索引8,DLPack contributors,DLPack: Open in-memory tensor structure for sharing tensors,2017–,https://github.com/dmlc/dlpack】的互操作性、一个稳定的C语言插件ABI,以及用于集成Triton【索引9 ,Triton contributors,Triton: A language and compiler for custom deep learning primitives,2019–,https://github.com/triton-lang/triton】和CUTLASS【索引10,NVIDIA ,CUTLASS: Cuda templates for linear algebra subroutines,2017–,https://github.com/NVIDIA/cutlass】等CUDA模板库编写的自定义内核的钩子 。
* AI辅助开发方法论:总结了一套实用的工作流程,用于通过智能体生成和验证系统级软件,该流程利用构建、测试、差异检查和多智能体代码审查作为保障机制。
* 产物评估:报告了代码库的规模和测试套件的构成,总结了附带的AI生成内核套件的可复现微基准测试(包括与PyTorch SDPA/FlashAttention对比的注意力机制),并报告了端到端的训练健全性检查(序列反转、CIFAR-10 ViT和miniGPT)以及一个仅在Blackwell上进行的多GPU扩展基准测试。
* 经验与局限性:描述了在生成系统软件中观察到的失败模式,包括一种“弗兰肯斯坦”(Frankenstein)组合效应,即局部正确的子系统在组合后可能导致全局性能不佳。

A3 背景知识与设计原则

背景与相关工作

即时执行与“定义即运行”。诸如Chainer【索引11,Seiya Tokui, Kenta Oono, Shohei Hido, and Justin Clayton,Chainer: a next-generation open source framework for deep learning,2015,Proceedings of Workshop on Machine Learning Systems (LearningSys) in The Twenty-ninth Annual Conference on Neural Information Processing Systems (NIPS),http://learningsys.org/papers/LearningSys_2015_paper_33.pdf】和DyNet【索引12 ,Graham Neubig et al,Dynet: The dynamic neural network toolkit,2017,arXiv e-prints,https://arxiv.org/abs/1701.03980】等“定义即运行”(define-by-run)系统确立了一种构建动态计算图的命令式方法。PyTorch【索引5 ,Adam Paszke, et al.,Pytorch: An imperative style, high-performance deep learning library,2019,arXiv preprint arXiv:1912.01703,https://arxiv.org/abs/1912.01703】通过其Python优先的接口和面向性能的C++核心推广了这一模型。VIBETENSOR遵循相同的即时执行 (eager execution)理念,但其重点在于验证使用编码智能体生成一个跨越多层(前端、运行时、CUDA分配器/图)软件栈的可行性。

内核编写系统。Triton【索引9,Triton contributors,Triton: A language and compiler for custom deep learning primitives,2019–,https://github.com/triton-lang/triton】等现代系统使得在Python原生工作流中编写自定义内核变得容易。而像CUTLASS【索引10,NVIDIA ,CUTLASS: Cuda templates for linear algebra subroutines,2017–,https://github.com/NVIDIA/cutlass】这样的CUDA模板库则为手写优化的自定义内核提供了补充方法。在本文中,CUTLASS指的是C++模板库;CuTeDSL则指我们附带的内核套件中使用的独立DSL后端。VIBETENSOR包含了对这两种风格的集成钩子,并提供了一个运用这些路径的内核套件 。

AI辅助软件工程。专注于代码的大型语言模型(LLMs)【索引6,Mark Chen et al.,Evaluating large language models trained on code,2021,arXiv preprint arXiv:2107.03374,https://arxiv.org/abs/2107.03374】和使用工具的智能体【索引7 ,Shunyu Yao et al.,React: Synergizing reasoning and acting in language models,2022,arXiv preprint arXiv:2210.03629,https://arxiv.org/abs/2210.03629】催生了新的工作流程,即通过迭代地提出编辑并利用工具进行验证来构建软件产物。像SWE-bench【索引13 ,Carlos E. Jimenez et al.,Swe-bench: Can language models resolve real-world github issues?,2023,arXiv preprint arXiv:2310.06770,https://arxiv.org/abs/2310.06770】这样的基准测试评估了在真实代码库中解决问题的能力。VIBETENSOR的不同之处在于,其产物横跨系统软件的多个层次,直至CUDA内存管理,并且其开源版本旨在作为一个研究基底,用于探究AI辅助工作流在系统规模下的行为 。

设计原则

VIBETENSOR针对的是设计空间中的一个务实子集。以下原则不仅塑造了其架构,也构成了生成过程中使用的评估标准。

带有显式控制的即时执行。VIBETENSOR遵循PyTorch【索引5,Adam Paszke, et al.,Pytorch: An imperative style, high-performance deep learning library,2019,arXiv preprint arXiv:1912.01703,https://arxiv.org/abs/1912.01703】推广的“定义即运行”模型:算子调用立即执行,自动求导引擎追踪一个动态图。这种模型非常适合调试、快速迭代和动态控制流 。

端到端的一致性。该系统在生成时就旨在贯穿从面向用户的API到GPU执行的完整路径。我们优先考虑实现张量、分发、自动求导和CUDA运行时机制的可用版本,即使它们尚不完整或未经过优化。

可测试性作为首要约束。在一个生成的代码库中,测试扮演着可执行规范和回归防护的角色。VIBETENSOR包含了C++ (CTest)和Python (pytest)测试套件,以及一个针对特定清单的、仅检查导入的API与PyTorch的对等性检查器。其Python CUDA接口包含了经过契约测试的运行时和分配器API(例如,memory_statsmemory_snapshot和每个进程的内存比例上限)。我们并不声称实现了全自动的测试生成;相反,这些测试被包含在发布的产物中,可供审查和重新运行。在该项目中,测试与实现是在同一个智能体循环中编写的:当发现故障时,工作流通常会通过添加一个最小化的回归测试或契约检查,然后重新运行测试套件来响应。这仍然是手写的测试代码(而非形式化生成),随着API接口的增长,测试的编写和维护仍然是一个扩展性成本。

互操作性与可扩展性。我们将该运行时视为一个更大生态系统中的一个组件。VIBETENSOR支持通过DLPack【索引8,DLPack contributors,DLPack: Open in-memory tensor structure for sharing tensors,2017–,https://github.com/dmlc/dlpack】进行零拷贝数据交换,提供了一个用于插件 的C ABI,并暴露了Python级别的覆盖机制以支持快速实验。

可观察性。生成的系统能够被观察是其一大优势。因此,VIBETENSOR在其CUDA缓存分配器中直接实现了分配器诊断功能——通过Python接口暴露的设备级统计数据和段/池快照API(例如,memory_statsmemory_snapshot)——以及CUDA图的检测工具和可选的多GPU可观察性钩子(见第4.7节)。这些诊断功能是VIBETENSOR原生的:其API受PyTorch启发,但计数器和快照反映的是VIBETENSOR自己的分配器状态,而不是调用PyTorch。

A2 方法细节

图1展示了VIBETENSOR的宏观架构。它被组织为前端、共享的核心运行时、CUDA子系统以及可选的内核和插件层。

前端:Python和Node.js

Python前端。VIBETENSOR提供了一个类似torch的Python覆盖层(vibetensor.torch),该覆盖层是基于一个nanobind【索引1,Wenzel Jakob and nanobind contributors,nanobind: Small and efficient bindings for python,2022–,https://github.com/wjakob/nanobind】扩展模块实现的。这个覆盖层提供了张量工厂函数、一个路由到调度器的操作命名空间 (ops namespace),以及用于流、事件、分配器统计和CUDA图捕获的CUDA实用工具。

Node.js前端。VIBETENSOR还包含一个基于Node-API【索引14,Node.js contributors,Node-api,2018–,https://nodejs.org/api/n-api.html】的实验性Node.js插件。这 个JavaScript/TypeScript (JS/TS)覆盖层是异步优先的:繁重的工作通过napi_async_work进行调度,以避免阻塞Node事件循环,并且通过一个进程范围的在途操作上限(VBT_NODE_MAX_INFLIGHT_OPS)来限制排队的工作量。在当前原型中,JavaScript张量对象仅支持CPU;CUDA功能则通过显式的H2D/D2H(主机到设备/设备到主机)传输和DLPack导入/导出功能来暴露。


图1:VIBETENSOR的高层架构:Python (nanobind) 和 Node.js (N-API) 前端将调用分派到一个共享的C++核心,该核心实现了张量/存储、分派、自动求导、索引、随机数生成器(RNG)以及CUDA运行时组件(流/事件/图和缓存分配器)。可选的内核库和动态加载的插件扩展了操作符集合。

核心运行时:张量、存储和视图

张量实现。C++核心将TensorImpl定义为一个对引用计数的Storage的视图,包含了尺寸/步长、存储偏移量、数据类型和设备元数据。这种设计支持非连续视图和as_strided语义,并维护一个(由所有视图共享的)原子版本计数器,以支持原地操作的安全性检查。

张量迭代器。对于逐元素和归约操作,VIBETENSOR提供了一个TensorIterator子系统,该子系统负责计算迭代形状和每个操作数的步长元数据。TensorIterator也通过插件ABI暴露,以便外部内核可以安全地重用迭代逻辑。

调度器:模式精简的算子注册表

调度器实现。VIBETENSOR实现了一个模式精简(schema-lite)的调度器,它将完全限定的算子名称(例如,vt::add)映射到内核实现。该调度器支持CPU和CUDA分派键、装箱(boxed)和非装箱(unboxed)调用路径,以及包装层(例如,自动求导和Python覆盖)。注册过程会发布每个算子的不可变快照状态,该状态编码了基础内核和包装层的存在情况,从而在稳定状态下实现无锁的调用路径。设备策略规则强制执行通用不变量(例如“所有张量输入都在同一设备上”),同时允许为实验性的多设备路径制定专门的策略。

反向模式自动求导

自动求导引擎。自动求导(Autograd)是作为一个反向模式引擎实现的,它包含Node/Edge图对象和每个张量的AutogradMeta。在反向传播过程中,该引擎维护依赖计数、输入缓冲区和一个就绪队列。对于CUDA张量,引擎会记录并等待CUDA事件,以同步跨流的梯度流。VIBETENSOR还包含一个实验性的多设备模式,旨在用于跨设备图执行的研究。

CUDA子系统:流、分配器和图

CUDA组件。VIBETENSOR为CUDA流和事件提供了C++封装,一个具有流序(stream-ordered)语义的缓存分配器,以及CUDA图的捕获/重放支持【索引2,NVIDIA,CUDA Graphs,2020–,https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__GRAPH.html】。该分配器包含诊断功能(快照、统计数据、比例上限和GC阶梯),以使内存行为在测试和调试期间可观察。CUDA图与分配器的“图池” (graph pools)集成,以管理内存生命周期跨越捕获和重放阶段。

互操作性与扩展

数据交换。VIBETENSOR通过DLPack导入/导出【索引8,DLPack contributors,DLPack: Open in-memory tensor structure for sharing tensors,2017–,https://github.com/dmlc/dlpack】支持CPU和CUDA张量的互操作性。它还包括一个用于高效序列化【索引15 ,Hugging Face and safetensors contributors,Safetensors: Safe and fast tensor serialization,2022–,https://github.com/huggingface/safetensors】 的C++20 safetensors加载器/保存器接口。

扩展机制。为了实现可扩展性,VIBETENSOR提供了:(i) 受torch.library启发的Python级覆盖机制,(ii) 一个用于动态加载插件的稳定C ABI,这些插件可以注册新的算子和内核,以及 (iii) 用于集成领域特定系统(如Triton【索引9,Triton contributors,Triton: A language and compiler for custom deep learning primitives,2019–,https://github.com/triton-lang/triton】或CUDA模板库如CUTLASS【索引10,NVIDIA ,CUTLASS: Cuda templates for linear algebra subroutines,2017–,https://github.com/NVIDIA/cutlass】)编写的自定义GPU内核的钩子。C插件ABI是版本化的,并暴露了基于DLPack的数据类型/设备元数据 和TensorIterator辅助函数,允许外部内核重用迭代逻辑并遵守别名规则。

Fabric与实验性多GPU组件

Fabric子系统。VIBETENSOR包含一个实验性的Fabric子系统,该子系统暴露了显式的点对点(peer-to-peer)GPU访问路径(当设备拓扑支持时,使用CUDA P2P和UVA)和可观察性(统计数据和事件快照)。其目的是支持对单进程多GPU执行的研究,在这种场景下,运行时可以对拓扑、同步和内存布局进行推理。在当前版本中,Fabric专注于显式的逐元素操作和可观察性原语,而不是一个完整的分布式训练运行时。

多GPU插件示例。作为硬件专用扩展的一个参考示例,该版本还包含一个尽力而为(best-effort)的基于CUTLASS的环形allreduce插件,目标是NVIDIA Blackwell级GPU;这是一个说明性插件(而不是基于NCCL的直接替代品)。该插件直接绑定CUTLASS实验性的环形allreduce内核,不调用NCCL;作为比较,插件目录中包含一个独立的PyTorch torch.distributed (NCCL) 带宽基准测试。图2展示了其宏观(GPU间环形结构)和微观(warp专用的流水线)视图。


图2:作为一个示例插件提供的经warp优化的环形allreduce内核(Blackwell SM100/SM103)的宏观和微观视图。该插件是尽力而为的,受限于工具链和硬件的可用性,并且不是使用VIBETENSOR核心运行时所必需的;它旨在作为一个参考实现,而非NCCL的替代品。

A7 补充细节

AI辅助开发方法论

VIBETENSOR是在高级别人类指导下,由LLM驱动的编码智能体在大约两个月内开发完成的。由于本文的重点是系统软件而非智能体实现,我们将智能体视为黑盒,它们提出代码更改并使用工具进行验证。人类提供高层次的需求和优先级,但不进行差异级别的审查或运行验证命令;相反,智能体执行构建、测试和差异检查,并在这些检查通过时保留更改。

开发工作流与保障措施。开发工作流反复应用一个简单的循环:(1)指定一个范围明确的目标和不变量,(2)生成并应用代码更改,(3)编译并运行专注的测试,以及(4)随着子系统的组合扩大验证范围。

关键保障措施。两个保障措施尤为重要:将测试作为规范(CTest/pytest,外加有针对性的多步回归测试),以及与参考实现进行差异检查(例如,针对选定算子与PyTorch比较、DLPack往返测试以及内核间的比较)。我们还使用多智能体代码审查来捕捉单个智能体可能忽略的缺失情况、冗余抽象和不安全模式。

案例研究:跨层调试。在智能体生成的系统软件中,一个反复出现的模式是,“单次”正确性(即一个算子一次性与参考实现匹配)并不意味着当该子系统在训练循环中反复组合时能够保持稳定。一次注意力内核移植过程的开发日志揭示了出现的各种问题:启动时限制(内核参数数量、共享内存使用量)、数值上的细微差别(稳定化注意力中的对数底和缩放),以及运行时风险,例如未初始化的GPU缓冲区,其重用问题仅在多次迭代后才会显现。

下表总结了代表性问题以及防止其再次发生的保障措施。


表1:在代理辅助开发期间观察到的代表性调试事件(选自开发笔记)。

局限性与经验教训

VIBETENSOR是一个研究原型,存在显著的局限性。

“弗兰肯斯坦”组合效应。在生成的系统中,一个反复出现的失败模式是,单个看似合理的组件组合在一起,可能形成一个全局次优的设计。在VIBETENSOR中,一个以正确性为优先的自动求导和分发设计可能会引入序列化点,从而使本应高效的后端内核处于“饥饿”状态。一个例子是用于拒绝并发多线程反向传播的非可重入全局反向门(实现为一个进程范围的try-locked互斥锁)。它简化了安全属性(例如,避免了嵌套反向传播中的死锁模式),但可能使独立的并行反向传播工作序列化。图6详细展示了这一瓶颈:首先,引擎中的一个全局门将并发的反向传播调用汇集起来;然后,主机端的序列化和同步传播到设备端;最后,本应高效的后端内核由于“饥饿”而利用率降低。


图6:“弗兰肯斯坦”效应的一个例子:一个高开销的前端/引擎层会使执行序列化,并导致高性能后端“饥饿”。这反映了当全局性能目标未在早期编码时,以正确性为先的生成方式的局限性。

不完整的API接口和性能。VIBETENSOR有意不追求与PyTorch完全兼容。许多算子、数据类型和分布式功能缺失或不完整,并且性能尚未经过调优以匹配生产级框架。

生成代码特有的验证差距。智能体生成的代码可能通过局部单元测试,但在重复组合(例如,多步训练循环)下因状态交互、未初始化的缓冲区或意外的全局同步而失败。这促使我们需要设计能够检验重复执行、跨流行为和长时间运行工作负载的回归测试。

维护、安全性和保障。机器生成的代码可能包含不一致的约定、冗余的抽象以及微妙的正确性或安全问题。因此,我们提醒不要在生产环境中使用,并将VIBETENSOR主要定位为一个研究和教育性质的产物。

A4 实验环境

A4 实验结果

我们从四个维度评估VIBETENSOR:代码库规模、正确性基础设施、选定组件的内核级性能以及端到端训练的健全性检查。

代码库规模

下表总结了本次发布的代码库规模,计数不包括供应商提供的第三方依赖。


表2:VIBETENSOR代码库规模(不包括供应商提供的第三方代码)。“核心”指C++运行时的include/和src/目录。

正确性与可复现性

VIBETENSOR包含C++单元测试(通过CTest运行)和Python测试(通过pytest运行)。测试覆盖了核心张量语义、调度器行为、自动求导、CUDA运行时工具、分配器不变量、DLPack互操作性、CUDA图和选定的内核路径。分配器诊断功能通过Python CUDA模块(vibetensor.torch.cuda)暴露,包括memory_statsmemory_snapshot和进程级内存比例上限;Python测试套件包含了对这些API的契约检查。为减少“能导入但功能损坏”的失败,VIBETENSOR还提供了一个仅导入的API对等性门控:一个脚本,根据一个范围受限的清单检查可导入符号与PyTorch的对应关系。

一个典型的验证序列如下:

python -m pip install -U pip build pytest numpy
CMAKE_BUILD_TYPE=Debug python -m pip install -v -e .[test]
ctest --test-dir build-py -j$(nproc) --output-on-failure
pytest -q
python tools/check_api_parity.py --manifest api/manifest_import_only.json
# (可选) Node.js覆盖层
cd js/vibetensor && npm ci && npm test

内核套件与微基准测试

附带的AI生成内核套件为选定的算子(如范数、旋转位置编码、优化器和softmax/损失函数)提供了Triton和CuTeDSL实现,以及记录性能和与PyTorch基线数值差异的基准测试工具。大多数Triton和CuTeDSL内核通过一个面向PyTorch的测试工具进行调用,该工具也作为差异检查的参考。对于部分内核,我们还提供了无需导入PyTorch即可运行的VIBETENSOR原生封装(通常通过DLPack)。此外,还包含一个专门的融合注意力基准测试,将一个Triton内核与PyTorch的缩放点积注意力(SDPA)进行比较。


图3:附带的AI生成内核套件的宏观视图。多个后端(Triton、CuTeDSL和PyTorch参考路径)共享一个面向Python的通用接口和基准测试工具。

下表报告了从内核套件报告和融合注意力基准测试中提取的部分微基准测试结果。所有计时都在测量区域前后进行同步,以反映设备时间而非主机分派延迟。


表3:从附带的内核套件和H100 PCIe上的融合注意力基准测试中提取的选定微基准(记录于发布产物中)。Torch指测试工具使用的PyTorch基线。Best报告该行中非Torch后端(如Triton, CuTeDSL)测得的最低延迟。这些结果描述的是孤立内核的性能,而非端到端模型性能。

注意力机制讨论:注意力基准测试结果好坏参半。对于一个NanoChat风格的训练配置(批量大小32,序列长度2048),Triton内核的性能优于PyTorch SDPA/FlashAttention基线(前向1.54倍,后向1.26倍)。然而,对于小批量GQA预填充工作负载,同样的基于指针的内核可能落后于FlashAttention(前向0.67倍,后向0.66倍)。这表明性能可移植性依赖于内核的特化程度和工具链的成熟度。

端到端训练运行

我们通过运行三个小型工作负载来验证VIBETENSOR能够组合成完整的训练循环:一个合成的序列反转任务、一个CIFAR-10视觉Transformer和一个在莎士比亚数据集上的miniGPT风格语言模型。这些工作负载在NVIDIA H100和Blackwell GPU上都进行了运行。


图4:在Hopper H100上比较VIBETENSOR和PyTorch在三个工作负载上的端到端训练曲线。在每种情况下,VIBETENSOR都匹配了PyTorch基线的定性收敛行为,表明核心张量语义、自动求导和优化器更新在训练循环规模上能够正确组合。

下表总结了两种架构下的平均计时。在这些工作负载中,VIBETENSOR目前比PyTorch慢(1.7–6.2倍),这与其原型状态和已知的序列化点相符。这些结果旨在作为端到端的功能性检查,而非性能声明。


表4:端到端训练工作负载和平均时间(记录于发布产物中)。迭代时间在注明时排除了第一个记录的步骤;epoch时间是在比较范围内平均的。

多GPU扩展(Fabric + 环形allreduce后端)

我们还在Blackwell GPU上报告了一个小型的多GPU基准测试,该测试使用了实验性的Fabric子系统(第4.7节)和一个环形allreduce后端(fabric_ring)。下图显示了代表性的训练曲线,下表则展示了在弱扩展(固定每个GPU的批量大小)情况下扩展到四个GPU的吞吐量。由于该路径依赖一个仅限Blackwell的CUTLASS插件,因此未在Hopper系统上运行。Fabric并非一个完整的分布式训练运行时;这些结果主要表明跨设备同步和可观察性路径能够端到端执行。


图5:Blackwell扩展基准测试的多GPU训练曲线(全局规模1–4)。


表5:在Blackwell GPU上使用VIBETENSOR Fabric的多GPU训练基准(弱扩展)。

A5 结论

VIBETENSOR项目证明了AI辅助工作流能够生成一个非平凡的、支持GPU的深度学习系统软件栈,该软件栈横跨Python和Node.js前端以及C++20/CUDA核心,并通过构建和测试进行验证。其开源发布旨在支持对系统级软件生成的可复现研究:哪些部分可以快速生成,会出现哪些类型的错误和瓶颈,以及什么样的测试和架构性支架能最有效地约束搜索空间。

从生态系统的角度来看,发布此类产物可以通过提供具体的、可修改的运行时和内核组件实现,并促使基于源代码和测试对AI辅助工程方法进行严谨的讨论,从而使更广泛的GPU软件社区受益。