人工智能机器学习深度学习【免费下载链接】frameworkA machine learning library with a TypeScript API and Rust backend. CUDA and WebGPU compatibility. Built to understand how ML frameworks and models work internally.项目地址https://gitcode.com/gh_mirrors/framework235/framework点击查看免费下载mni-ml/framework是一个 TypeScript API Rust 原生的机器学习框架内置自动求导、张量运算与 GPU 加速CPU / CUDA / WebGPU 三后端。如果你想给它新增一个算子本文带你走完 TypeScript → Rust → CUDA/WGSL 的完整链路只需看懂 5 个文件就能学会整个框架的算子扩展方法。一、先搞懂架构一个算子为什么会穿过 5 层框架的调用链非常清晰TypeScript API (src/tensor.ts) │ 通过 N-API 调用 ▼ N-API 桥接层 (src/native/src/lib.rs) │ 转发给具体算子模块 ▼ Rust 算子实现 (src/native/src/ops/*.rs) ├── cpu / webgpu 共用纯 Rust 实现 └── cuda 分支启动 GPU 内核 ├── CUDA 内核 (src/native/kernels/*.cu) └── WebGPU 着色器 (src/native/shaders/*.wgsl)同时每个算子都要在 src/native/src/autograd.rs 的autograd tape反向传播磁带中登记这样backward()才能找到它的反向公式。层文件职责TS APIsrc/tensor.ts给Tensor类添加用户可调用的方法N-APIsrc/native/src/lib.rs#[napi]函数连接 JS 与 RustAutogradsrc/native/src/autograd.rs定义BackwardOp、登记反向函数算子核心src/native/src/ops/前向/反向的 CPU 与 CUDA 实现GPU 加速src/native/kernels/ / src/native/shaders/CUDA 内核与 WGSL 着色器 新手提示三个后端由 Cargo feature 在编译期互斥选择cpu/cuda/webgpu见 src/native/Cargo.toml。所以每个算子通常要写 2 套实现cpuwebgpu共用一套纯 Rust 代码cuda单独一套启动 GPU 内核。二、以新增mygelu激活函数为例5 步走完全流程假设我们想新增一个自定义激活函数以类似 GELU 的mygelu为例下面是完整的修改清单。步骤 1TypeScript 层——给 Tensor 暴露方法在 src/tensor.ts 中找到激活函数区域仿照relu()添加一行转发方法。核心套路只有一句把张量 id 传给 native拿回新张量 id// src/tensor.ts mygelu(): Tensor { return new Tensor(native.mygelu(this._id)); }TS 层完全不需要关心计算细节——所有张量数据都存放在 Rust 侧的TensorStore中_id就是句柄。步骤 2N-API 桥接层——在 lib.rs 注册入口在 src/native/src/lib.rs 的 Ops 区域新增一个#[napi]函数。它负责加锁拿到全局Engine再把store和tape传给算子实现// src/native/src/lib.rs #[napi] pub fn mygelu(a: u32) - u32 { let mut e engine().lock(); let Engine { store, tape, .. } mut *e; ops::activation::mygelu_forward(a as TensorId, store, tape) as u32 }⚠️ 命名细节Rust 函数名用下划线mygeluN-API 会自动转成 JS 驼峰名所以 TS 侧调用native.mygelu(...)即可。步骤 3Autograd 登记——新增 BackwardOp这是最容易被新手漏掉的一步。打开 src/native/src/autograd.rs做两件事在BackwardOp枚举中加一个变体MyGelu在底部的dispatch_backward匹配中加一条路由指向你的反向函数// src/native/src/autograd.rs BackwardOp::MyGelu ops::activation::mygelu_backward(grad_id, entry.saved, store),SavedContext如SavedContext::Tensor(a)决定前向阶段要为反向保存哪些中间量比如 GELU 就保存了输入张量见 src/native/src/autograd.rs。步骤 4Rust 算子核心——前向 反向 × 2 套后端以 src/native/src/ops/activation.rs 中的 GELU 为模板需要写 4 个函数函数编译条件做什么mygelu_forward#[cfg(any(feature cpu, feature webgpu))]纯 Rust 逐元素计算结果写入 storemygelu_forward#[cfg(feature cuda)]启动 GPU 内核指针直接操作显存mygelu_backward同上两套用SavedContext中的中间量计算梯度CPU/WebGPU 版前向的核心模式注意最后必须tape.record登记磁带#[cfg(any(feature cpu, feature webgpu))] pub fn mygelu_forward(a: TensorId, store: mut TensorStore, tape: mut Tape) - TensorId { let data: Vecf32 store.to_host(a).iter().map(|x| mygelu_scalar(x)).collect(); let shape store.shape(a).to_vec(); let out store.from_vec(data, shape); tape.record(TapeEntry { op: BackwardOp::MyGelu, output_id: out, input_ids: smallvec![a], saved: SavedContext::Tensor(a), }); out }CUDA 版则是取指针 → 启动内核 → 同样tape.record#[cfg(feature cuda)] pub fn mygelu_forward(a: TensorId, store: mut TensorStore, tape: mut Tape) - TensorId { let shape store.shape(a).to_vec(); let n shape_size(shape); let out store.zeros(shape); let dev GpuDevice::instance(); let func dev.get_func(mygelu_forward_f32); unsafe { dev.stream.launch_builder(func) .arg(store.dev_ptr(out)) .arg(store.dev_ptr(a)) .arg((n as i32)) .launch(launch_cfg(n as u32)) .unwrap(); } // ... tape.record(...) 同上 out }反向推导小技巧激活函数的反向几乎都是dx dy * f(x)。GELU 的梯度推导可以参照 src/native/src/ops/activation.rs 中的gelu_grad_scalar。步骤 5GPU 加速——CUDA 内核与 WGSL 着色器CUDA 侧在 src/native/kernels/elementwise.cu 中添加一个__global__内核签名与exp_f32保持一致输出指针、输入指针、元素个数extern C __global__ void mygelu_forward_f32(float* out, const float* a, int n) { int i blockIdx.x * blockDim.x threadIdx.x; if (i n) out[i] mygelu_kernel(a[i]); }然后在 src/native/src/device.rs 的load_all_kernels中把内核名注册进对应列表否则运行时会报Kernel function mygelu_forward_f32 not found。WebGPU 侧在 src/native/shaders/elementwise.wgsl 中加一个 compute 入口绑定约定是binding(0)out, binding(1)a, binding(3)params.ncompute workgroup_size(256) fn mygelu_forward_f32(builtin(global_invocation_id) gid: vec3u) { let i gid.x; if (i params.n) { out[i] mygelu_wgsl(a[i]); } }同样在 device.rs 的load_all_shaders中登记入口名即可。✅ 好消息CUDA 与 WebGPU 内核逻辑完全一致只是语法不同。逐元素算子两者都是一个线程处理一个元素 越界判断照抄exp_f32即可。三、构建与测试如何验证你的新算子一键构建各后端package.json 已配好三个构建脚本任选其一# CPU默认无需 GPU npm run build:native # CUDA需 CUDA toolkit npm run build:native:cuda # WebGPU npm run build:native:webgpu它们本质是在src/native下执行cargo build --release并拷贝产物为.node文件。用梯度检验反向公式框架自带轻量测试框架参考 test/autograd.test.ts 的写法即可验证数值梯度以x²在x3处导数应为 6 为例const x Tensor.fromFloat32(new Float32Array([3]), [1]).setRequiresGrad(true); const y x.mygelu(); y.backward(); assertClose(x.grad!.toFloat32()[0], expected, 1e-3, mygelu grad);运行npm test即可执行全部测试套件见 test/run.ts。测试技巧取几个解析解已知的点如 0、±1对比手算导数是最快的反向公式校验方法。四、新手常见坑位清单 坑症状解决忘了在dispatch_backward登记backward()时梯度静默丢失补上 autograd.rs 的 match 分支CUDA 内核没注册Kernel function xxx not found在 device.rs 的load_all_kernels加名字忘了tape.record前向结果正确但反向报错/梯度为 0每个 forward 末尾都要 record只改了 CPU 分支CUDA 构建时编译报错或行为不一致注意#[cfg(feature)]分支要成对写TS 方法名与 N-API 导出名不匹配native.xxx is not a functionRust 下划线名会被自动转驼峰注意对应五、总结算子扩展完整路线图TS 暴露src/tensor.ts 加方法一行转发N-API 注册src/native/src/lib.rs 加#[napi]函数Autograd 登记autograd.rs 加BackwardOp变体 dispatch 分支Rust 实现ops/ 下写前向/反向 × cpu/webgpu cuda 两套GPU 内核kernels/*.cu 与 shaders/*.wgsl 各加一个入口并在 device.rs 注册构建测试npm run build:nativenpm test验证梯度这套TS 句柄 Rust 计算 编译期选后端 磁带自动求导的设计正是理解各类 ML 框架内部运作的绝佳样本——读懂一个add或gelu的完整链路推荐从 src/native/src/ops/elementwise.rs 读起你就已经掌握了框架 90% 的扩展能力。赞分享人工智能机器学习深度学习【免费下载链接】frameworkA machine learning library with a TypeScript API and Rust backend. CUDA and WebGPU compatibility. Built to understand how ML frameworks and models work internally.项目地址https://gitcode.com/gh_mirrors/framework235/framework点击查看免费下载相关推荐在浏览器中跑深度学习mni-ml/framework 的 WebGPU 后端与 WGSL 计算着色器完全指南在浏览器中跑深度学习mni ml/framework 的 WebGPU 后端与 WGSL 计算着色器完全指南 mni ml/framework 是一个用人工智能机器学习深度学习Rust 原生二进制的跨平台分发mni-ml/framework 的 npm optionalDependencies 方案全解析Rust 原生二进制的跨平台分发mni ml/framework 的 npm optionalDependencies 方案全解析 想让 TypeScrip人工智能机器学习深度学习一份代码三种后端mni-ml/framework如何实现CPU、CUDA与WebGPU统一架构一份代码三种后端mni ml/framework如何实现CPU、CUDA与WebGPU统一架构 mni ml/framework 是一个 TypeScri人工智能机器学习深度学习创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
