我今年18岁。没有机器学习背景,没有上过大学,高中都没去上,没有导师。
几天前我盯着一张纸发呆。突然想:为什么计算机神经网络一定要是2D的?可以模拟生物吗?为什么一定要在平面上算?如果多个平面,岂不是翻倍?如果把六张纸想象成一个魔方,六个面各自承载神经元,八条体对角线变成新的通信通道会怎么样?
我真的很喜欢折腾这些,然后我立刻制定了详细计划,使用AI工具辅助写下了第一个 kernel。跑崩了。我又重新想了一下,和qq群友分享了我的目标,又写。又崩。连续几十次。没有 PyTorch,没有 TensorFlow,没有 CUDA。只有Swift和Metal。因为我的电脑显卡是AMD Vega 64,没装任何框架辅助,因为我想明白最底层的运行方式是什么原理。
这就是CubeNN。
##以下为AI的详细解答,内容与架构改动太多,我在这里一次讲不清楚
它是什么
一个用魔方几何作为计算架构的神经网络引擎。
标准 Transformer: 把数据排成一行,O(n²) 地互相看
CubeNN: 把数据分布在 14 个面上,只在该看的地方看
6 个标准面 → 块稀疏注意力(粗看全局 + 细看局部)
8 个 X 面对角线 → 跨面信息桥(不做 Attention,只负责传递)
每轮:6 面算 → 投影到 8 X 面 → 上采样精炼 → 融合回 6 面
最关键的是 Cube Cascade——一个树+链级联推理:
树阶段: 1 个魔方 spawn 8 个 → 8 个 spawn 64 个 → 73 个并行探索
GPU 上同时跑,选最优路径
链阶段: 最优叶子无限深度精炼
3-5 步收敛,方差提升 ~7%
怎么实现的
纯 Swift + Metal。零依赖。零框架。
// 大致代码就是这些
import Metal
import Foundation
let device = MTLCreateSystemDefaultDevice()!
let library = try! device.makeLibrary(filepath: "cube_nn.metallib")
// ...12 个 GPU kernel,12,000 次 dispatch
关键技术决策:
单 Command Buffer:整个树阶段 73 个魔方的全部 kernel dispatch 打包进一个 CB,0 次 CPU-GPU 同步
Pipeline State 缓存:编码从 1022ms 降到 42ms
Buffer 偏移:所有 73 个魔方的 14 个面存进一个连续 buffer,kernel 通过 buffer(15) 传偏移量
FP16:N≥64 时半精度提速 21%
性能
##经过测试,但是因设备差异可能不准确,仅参考
AMD Radeon RX Vega 64 (2017 年显卡, 14nm, 295W):
规模
神经元
魔方数
耗时
N=32
6,144
73 (树)
435ms
N=64
24,576
21 (树)
817ms
N=128
98,304
1
116ms
N=32 全连接 Attention 每层 201M FLOP → CubeNN 块稀疏 370K FLOP (544× 减少)
N=128 全连接需要 32GB 显存(物理上不存在)→ CubeNN 用 192KB
N=256 全连接需要 2.2T FLOP → CubeNN 52M FLOP (42,300× 减少)
代码体积:161KB。 对比 PyTorch 的 800MB。
我经历了什么
这个项目最困难的不是写 kernel,是在没有任何人告诉我"能不能做"的情况下,靠反复试错找到路。
第一次试图跑 73 个魔方,GPU 直接 hang 了。花了 3 天定位到是 Command Buffer 堆叠过多。
改了 single encoder 方案,又碰上 SIGILL——Metal 不允许 makeBuffer(length: 0),B=0 时创建了零长度 buffer。
想用 threadgroup memory 做 kernel fusion,结果跨 threadgroup 读不到数据,才明白 LDS 是 per-group 的。
N=64 的 FP16 要手动写 float↔half 转换函数,因为 macOS 11 上 Float16 类型被标为 unavailable。
每一次崩溃都教会我一个 Metal 的底层细节。没有人教我,但 Metal 的报错信息就是最好的老师。
为什么发在 Apple 开发者论坛
因为这是为苹果生态而生的项目。CubeNN 从头到尾只用了两个东西:Swift 和 Metal。它不需要移植就能跑在任何 Apple Silicon Mac 上(API兼容)。如果未来能把部分 kernel 映射到 Neural Engine,效率会再翻几倍。
我想问 Apple 的 Metal 工程师和 Core ML 团队:
** 有没有更好的 GPU 任务调度方式?**目前表现仍然欠佳(对于我这个完美主义者来说),可能改得有点乱了
有没有兴趣评估这个架构在 M4 上的表现? 我手里只有 Vega 64。M4 GPU + ANE方法 跑 CubeNN 会是什么效果?
源代码
├── run.swift # 统一 CLI,参数化 N/B/depth
├── src/
│ ├── cube_nn.metal # FP16 kernel
│ └── cube_nn_fp32.metal # FP32 kernel
└── benchmarks/ # 实测数据
如果你读到了这里——谢谢你。一个门外汉靠痴狂的,纯粹到几乎是妄想的主意和Metal走到了这里。我懂的不是很多,如果这个架构有任何价值,我想让它变得更好。任何建议、批评、或者指教,都非常欢迎。
0
0
35