探索一个轻量级CUDA驱动的深度学习库 | 原创,AI翻译

Home 2026.05

两位参与者:azticelzwjava

以下对话使用 Whisper 在 RTX 4070 上转录,然后借助 AI 工具进行精炼和整理。原对话语言为中文。

注:由于转录和精炼均由 AI 辅助完成,部分细节可能不准确或经过改写。在依赖任何重要信息前,请自行核实。


本次会议围绕一个名为 Axono 的个人开源项目展开了广泛的技术讨论——这是一个轻量级、CUDA 加速的张量计算库,功能类似于极简版的 PyTorch。对话探讨了该库的架构、CUDA 内核实现、CPU 与 GPU 之间的内存管理,以及单人开发的挑战。还深入介绍了开发者参与竞争性 AI 工程的经验,包括在一次国际大数据竞赛(IKCEST 2025)中获得第 17 名的佳绩,以及正在参与的一个模型优化挑战赛,目前在 16 名决赛选手中排名第 8。对话中还涉及了 AI 辅助编码等工具实践、V2 重设计从纯 CUDA 转向 CUBLAS 的原因、参与者之间交流语言的选择,以及动手工程能力相对于理论知识的价值。

Axono 项目:一个极简类 PyTorch 的 CUDA 库

项目概述与动机

讨论从一位参与者介绍一个名为 Axono 的项目开始,该项目托管在一个 GitHub 组织中。开发者将其描述为一个轻量级推理库,并指出它支持多种数据精度类型,包括 int8、int16、int32、int64、float32 和 float64。该库包含常用算子:矩阵乘法(matmul)、加法以及简单的激活函数。当被问及它是否像 PyTorch 一样是 Python 库时,开发者确认了这一点,但强调由于只有一个人在开发,进展已大大放缓——这使得开发过程相当疲惫。

该库主要面向 Linux 系统,虽然用于演示的机器上存在 Windows 版本,但没有官方的 Windows 打包版本。开发者指出,当前展示的版本是 1.0,但澄清说演示的并非最新的 1.0 版本——而是一个较旧的发行版。当被问及已经发布了多少个版本时,回答是大约十几个到二十个。开发环境位于一台配备 NVIDIA V100 GPU 的机器上。

当前状态与 V2 重设计

对话转向了项目结构。开发者解释说,当前的代码库(约几千行)相当混乱,这促使他们为 V2 版本进行了完全重写。V0 版本存在严重的时间问题,因此 V1 正在从头开始开发以取代它。然而,V1 尚未完成,因为整个过程相当麻烦。一个关键的架构变化是:V1 将不再支持纯 CUDA——它将强制使用 CUBLAS。不过,开发者提到,他们最终可能会为那些仍需要纯 CUDA 的老旧设备提供兼容性。

团队讨论了一个问题:能否用 Axono 编写一个小型神经网络,比如用于 MNIST 的数字分类器?该库目前不支持线性层(Linear layer)——这是构建此类模型的关键需求。开发者承认了这一限制,但指出由于 Axono 具备矩阵乘法功能,理论上可以用数组来表示线性层。大家一致认为这将是一个有趣的实验,一位参与者表示他们会获取代码并尝试探究,尽管缺少 Linear 抽象,是否仍能实现一个极简的 MNIST 实现。

张量操作与形状处理

一位参与者演示了如何声明一个张量变量并探索其形状。在展示矩阵乘法时,他们举例说明:矩阵 A 的形状为 (2, 3),矩阵 B 的形状为 (3, 2),结果形状应为 (2, 2)。关于广播机制与标准矩阵乘法之间存在一些混淆——开发者澄清说,该库并不直接融合形状,而是将内维度相乘,外维度作为结果形状输出。

关于形状变换,开发者表示目前只提供了转置方法。没有通用的 reshape 或形状变更方法。转置实现只是简单地翻转行和列(A[i][j] 变成 A[j][i]),但相关代码相对冗长。当被问及更高维度的张量时,该库仅支持二维转置——它不能同时处理多个维度的转置。开发者指出,约定是显式传入两个维度(dim0 和 dim1)。

CUDA 实现:从 CPU 到 GPU 再返回

内核启动流水线

开发者从矩阵乘法开始解释了底层的 CUDA 实现。过程首先分配内存,设置启动配置,然后调用内核。在内核内部,通过一个简化的行列循环计算得到结果。该库包含优化措施:其中一个内核名为 matmul_optimized,采用列分块方法,旨在使整体计算运行得更快。

为了优化性能,开发者使用了块大小和共享内存。具体来说,他们将块大小参数放入一个临时模板中,将数据拉入共享内存块以加速矩阵乘法。分块概念涉及将每个矩阵分割成小块(tiles),以提高计算效率——类似于 Flash Attention 中使用的分块技术。

共享内存与内核结构

一位参与者要求在此上下文中解释共享内存。开发者坦言没有完美的直观解释,将其描述为一种将矩阵块加载到共享内存(GPU 上更快的存储空间)以加速重复访问的技术。代码中还包含边界检查,以确保索引不会超出源张量和目标张量的范围。idx 变量用于计算转置操作中源张量和目标张量的位置。

矩阵乘法代码的跟踪显示,该库在内核启动前会检查张量形状是否兼容。calculate_launch_config 函数计算网格大小和块大小。块大小设置为 256。网格大小公式采用“向上取整”方法:(num_elements + block_size - 1) / block_size。一位参与者起初对这个公式提出疑问,想知道为什么不用 num_elements 直接除以 block_size。开发者解释说,+ block_size - 1 的技巧是为了实现整数除法向上取整——当元素数量不能被整除时,该公式会向上取整,分配足够的块。这是竞赛编程(OI)和通用 GPU 内核设计中的常见技术。

CPU 与 GPU 之间的内存传输

对话深入探讨了数据如何在 CPU 和 GPU 之间移动。开发者描述说,在内核执行期间,数据被直接“扔进”内核——这就是传输过程。当被问及哪个库处理这种传输时,开发者表示是他们自己的自定义封装。然而,他们澄清说,数据实际上并不是在内核启动时从 CPU 移动到 GPU 的;张量从一开始就在 GPU 上初始化和分配。传递给内核的只是 GPU 内存的指针。然后内核直接对该 GPU 驻留数据进行操作。

开发者打了个比方:CPU 通过指针调用 GPU 函数,GPU 内核通过提供的指针访问该内存。他们强调,GPU 和 CPU 是分开运行的——CPU 不能直接访问 GPU 内存,因此 CUDA(或 CUBLAS)提供了调用 GPU 函数的机制。GPU 计算完成后,数据可以传回 CPU 进行显示或进一步处理。讨论中提到了术语“同步”(sync):有时 CPU 必须与 GPU 同步,等待 GPU 完成任务后再继续。在 PyTorch 中,许多过程默认是异步的,但显式的同步操作可确保数据完整性并避免错误。大家一致认为,操作体系主要由 CPU 驱动,GPU 结果必须传回 CPU 才能进行面向用户的输出或后续步骤。

代码走查:关键组件

Dispatch_Fill 模式

团队检查了代码库中的一个 dispatch_fill 函数。其目的是用特定值填充张量。它通过基于数据类型进行分发——不同数据类型映射到不同的内部实现,本质上在 C++ 中实现了一种泛型编程。该函数计算启动配置,然后将值填充到张量的数据指针中。涉及到 memory_copy 步骤:一个版本使用源指针和 fill_value 将数据从一个位置复制到另一个位置。开发者确认,tensor.data 只是一个原始指针,而不是一个库对象。

逐元素加法内核

开发者展示了一个逐元素加法内核。该内核执行逐元素加法:对于每个索引 idx,计算 result[idx] = a[idx] + b[idx]。这是一种直接的并行模式,每个线程处理一个元素。内核包含一个约束:仅当 idx 小于 num_elements 时才进行计算,作为防止越界访问的守卫。开发者指出,对于理解这类细节概念,他们通常更倾向于将代码复制到 AI 工具中并请求解释,而不是自己记住这些细节。

转置内核

详细检查了转置 CUDA 内核代码。该内核处理二维转置,参数包括源、目标、维度大小、源步长、目标步长以及维度索引(dim0, dim1)。计算使用 idx 来计算坐标,然后通过将坐标乘以步长并加上偏移量来计算源索引和目标索引。在索引计算过程中,使用 remainder 运算来计算剩余维度。开发者指出,该内核包含对 INF(无穷大)值的处理——当值超出可表示范围时,会生成 INF 而不是导致错误,这是可以接受的行为。

开发实践:AI 辅助与代码审查

开发者坦率地表示,大部分 CUDA 内核都是借助 AI 编写的——他们不会从头手写所有这些代码,因为那会太累。然而,他们强调自己清楚地理解每段代码的作用。使用 AI 仍然需要迭代:调整提示词、处理编译错误、调试以及确保正确性。他们通常要求 AI 每次生成 50 到 100 行代码,而不是一次性生成整个组件,然后修复出现的微小错误。一位参与者评论说,代码结构看起来非常精致,考虑到复杂度,开发者的方法令人印象深刻。

竞争性 AI 工程与模型优化

IKCEST 2025 竞赛:在 1700 支队伍中排名第 17

开发者分享了他们在 IKCEST 2025 国际大数据竞赛中的经历,他们在约 1700 支参赛队伍中获得了第 17 名。该竞赛要求队伍使用学校信息进行注册——队伍中必须至少有一名来自注册院校的在读学生。开发者以香港的一支队伍身份参赛,使用了所在学校的资质。

开发者展示了他们的竞赛 PPT。队伍名称为“everywhere we go”(原始幻灯片中存在一个小拼写错误)。开发者承认在香港很难找到可靠的队友;他们个人不认识其他有相似兴趣的人。对于后续的一场竞赛,他们与一名研究生组队,这名研究生能够提供有意义的帮助,不像典型的本科生那样。

竞赛工作流程与技术方法

该竞赛涉及一个问题,要求从图形输入中推导出答案。开发者的流水线第一步使用 OCR 识别,具体来说是 PPOCR(百度 PaddlePaddle 框架中的 PaddleOCR)。他们指出,在香港很多人不喜欢国内的国产工具,更倾向于西方替代品,但开发者使用 PaddleOCR 时没有这种偏见,并且发现它效果很好。OCR 结果产生了字符串提取,团队将其与一个检测网络(类似于 YOLO)结合起来,在识别之前定位文本区域。

竞赛期间的排行榜指标包括准确率(在某示例中为 0.47381),记录日期为 10 月 27 日。开发者指出具体指标细节有些遗忘,但代表了识别结果和字符串提取质量。竞赛中排名靠前的队伍使用了多种优化技巧来实现更快的推理速度,不过开发者记不起他们所有的具体技术。

当前竞赛:模型服务优化

开发者目前正在参加另一个关于模型服务优化的竞赛。该竞赛开始时约有 480 支队伍,经过一轮淘汰后迅速减少到 16 名决赛选手。开发者所在的队伍在这 16 名决赛选手中排名第 8,队伍名称为“forgive me out”。他们的队友是来自华东师范大学的一名研究生。

竞赛指标包括准确率(上限为 0.7——超过 0.7 的分数视为 1.0,大多数队伍很容易达到 0.7)、TTFT(首 Token 时间)和吞吐量(每秒 Token 数)。开发者认为 TTFT 和吞吐量才是真正困难的指标——准确率基本上是可以保证的。他们在一个月左右的时间里累计提交了 54 次,其中 27 次提交计入排名。他们的吞吐量得分是 38(基于竞赛使用的任何尺度)。排名第一的队伍在所有三项指标上均领先。

对于这次竞赛,开发者从头开始手工编写了解决方案,理由是 AI 生成的代码对于这类工作来说太容易出错。他们花了大约两周时间将模型编写成一个可部署的图表示。他们目前处于最终阶段,前三名队伍将前往越南参加最终比赛。

衡量工程能力:终极考验

当被问及在该领域什么是高水平的工程能力时,开发者给出了一个具体的基准:成功从头编写一个多模态语言模型(如 Qwen-VL-2B 或类似模型),并将其部署到 Kubernetes 集群中。这需要对注意力机制有深入理解,并且能够处理 CUDA 内核的约束——例如,不允许动态形状;一切都必须稳定。开发者表示,实现这一点表明技能水平大致与他们相当,并强调这确实很困难,需要多年的专注工作。

交流语言的选择

参与者在讨论过程中尝试切换为英语,但遇到了挑战。虽然两人都能用英语交流,但技术词汇(尤其是 CUDA 相关的术语)造成了理解困难。一位参与者指出,他们的口语英语词汇在会话场景中比在技术术语中更丰富,因此当遇到不熟悉的技术术语用英语表达时,他们有时会很吃力。大家一致认为,对于这类技术讨论,中文交流的效率大约是英语的两倍,同时承认提高英语水平对于国际平台和未来职业发展很有价值。他们决定本次对话主要用中文继续,并打算随着时间的推移多练习英语。

个人背景、未来计划与更广泛的兴趣

学习历程与技术广度

开发者已经从事这些系统相关工作大约三年。他们的知识不仅限于神经网络算法,还涵盖广泛的主题,包括多模态模型(VLM)、检测模型和 OCR 流水线。他们将自身的专业知识描述为“广而杂”,承认自己广泛探索而非在单一领域深入钻研。除了 Axono 库之外,他们的 GitHub 上还有更早的项目:一个基于 PHP 的小型项目、一个 JSON 数据库实验,以及一个刚学 Python 时创建的很早的项目。他们运营的 GitHub 组织是许久之前建立的;基本上是一个单人项目。

即将到来的行程

会议在五月举行。开发者提到他们当天晚些时候需要前往香港城市大学——他们目前身处深圳,计划过境。他们参观过许多香港大学校园(香港中文大学、香港大学、香港理工大学、香港科技大学),有时是去上课,有时是与朋友一起参加活动。当被问及这些学校学生的普遍技能水平时,他们拒绝给出直接评价,指出根据他们的经验,许多学生处于平均水平,在实战竞赛中可能无法胜过他们。

合作兴趣

一位参与者表达了合作参加竞赛的兴趣,认为自己的基础设施背景可以补充开发者的优势。开发者坦诚回应:他们认为参与者“目前还不够格”参与大数据基础设施竞赛所需水平,指出这类竞赛既要求扎实的系统知识,也要求算法技能。不过,他们给出了明确的路径:如果参与者能够成功从头编写一个多模态 LLM 并将其部署到 Kubernetes 集群中,那将证明具备同等能力。开发者表示,一旦达到那个水平,他们愿意合作。

结语

对话在覆盖了大量技术内容后结束。参与者们承认讨论了很多内容,其中一人建议暂时休息一下。开发者需要准备午餐然后前往香港,这提供了一个自然的结束点。会议在友好的氛围中结束,大家非正式地表示保持联系,并可能未来在边界对面见面。


项目链接:GitHub 上的 Axono


Back Donate