Kuma：把 PyTorch 模型编译成浏览器可执行文件

一个叫 Kuma 的开源项目最近在 Reddit 机器学习社区引发讨论。它的思路很直接：把导出的 PyTorch 模型编译成一个自包含的包，扔进浏览器就能跑，不需要 Python 环境，不需要服务端推理，甚至不依赖任何重量级运行时。

这事听起来简单，做起来不简单。

它到底在解决什么问题

传统的模型部署路径大概是这样：你在 PyTorch 里训练好模型，导出成 ONNX 或 TorchScript，然后要么部署到服务器跑推理 API，要么用 TensorRT、ONNX Runtime 这类运行时在边缘设备上跑。无论哪种，都绑定在特定的运行环境上。

浏览器端推理这几年也有不少尝试——TensorFlow.js、ONNX.js、Transformers.js，但它们本质上是把一个「运行时」塞进浏览器，模型文件是模型文件，运行时是运行时，分开加载。

Kuma 的做法不一样。它把模型编译成一个完全自包含的 artifact，包含：

• 计算图结构
• 二进制权重
• 后端 kernel（目前是 WGSL 格式）
• 运行时元数据

这个包本身就是可执行的。配合一个轻量级的 WebGPU 运行时，浏览器直接加载执行，中间没有额外的格式转换，没有动态图解释，没有 Python 依赖。

用作者 Slater Victoroff 的话说，这是「单一可移植 artifact」的理念。

为什么是 WebGPU

这是个好问题。WebGL 在浏览器里搞 GPU 加速已经很多年了，为什么 Kuma 选择 WebGPU？

简单说，WebGL 是为图形渲染设计的，要在上面跑通用计算（GPGPU）得各种 hack。矩阵乘法？你得把它伪装成纹理采样。batch normalization？又是一堆 shader 魔法。性能损耗不说，开发体验极差。

WebGPU 不一样。它从设计上就支持计算着色器（compute shader），API 更接近 Vulkan 和 Metal 的思路，能直接访问 GPU 的通用计算能力。更重要的是，它的着色器语言 WGSL 是专门为此设计的，不用再把计算逻辑硬塞进图形管线。

截至 2026 年，主流浏览器对 WebGPU 的支持已经相当完善——Chrome、Edge、Firefox、Safari 都能跑。这意味着 Kuma 编译出的包理论上在任何现代浏览器里都能执行，跨平台问题基本解决。

技术架构拆解

来看看 Kuma 的编译流程具体是怎么做的。

第一步：模型导出

起点是 PyTorch 的 torch.export 或 TorchScript。这一步把动态的 Python 代码转成静态的计算图表示。PyTorch 2.x 的 torch.compile 生态在这方面成熟了很多，TorchDynamo 能安全地捕获大部分 PyTorch 程序，AOTAutograd 处理反向传播的图捕获。

Kuma 在这基础上做进一步处理，把计算图转成自己的中间表示。

第二步：Kernel 生成

这是核心步骤。对于计算图中的每个算子（卷积、矩阵乘、激活函数等），Kuma 生成对应的 WGSL 计算着色器。

WGSL 代码大概长这样：

@group(0) @binding(0) var<storage, read> input: array<f32>;
@group(0) @binding(1) var<storage, read> weight: array<f32>;
@group(0) @binding(2) var<storage, read_write> output: array<f32>;

@compute @workgroup_size(256)
fn matmul(@builtin(global_invocation_id) gid: vec3<u32>) {
    let row = gid.x;
    let col = gid.y;
    var sum: f32 = 0.0;
    for (var k: u32 = 0u; k < K; k = k + 1u) {
        sum = sum + input[row * K + k] * weight[k * N + col];
    }
    output[row * N + col] = sum;
}

当然，实际生成的 kernel 会做更多优化——内存合并访问、共享内存利用、向量化等。这些优化直接影响最终性能。

第三步：权重打包

模型权重以二进制格式嵌入最终的 artifact。这里有个设计决策：是存原始 float32，还是量化成 float16 甚至 int8？

对于浏览器场景，下载体积很关键。一个 100MB 的权重文件在移动网络下加载时间可能超过 10 秒。Kuma 目前支持权重量化，但具体策略还在迭代中。

第四步：生成自包含包

最后，计算图、kernel、权重、运行时元数据打包成一个文件。浏览器端的运行时加载这个文件，按计算图拓扑顺序调度 kernel 执行。

整个流程是一次编译，到处运行——前提是目标环境支持 WebGPU。

实际跑起来是什么效果

目前 Kuma 仓库里的 demo 主要是神经视频表示（Neural Video Representation）。这类模型相对简单，输入坐标输出颜色值，适合验证端到端流程。

但作者的真正目标是算子网络（Operator Networks）和科学机器学习（Scientific ML）。这类应用的特点是：

1. 模型相对小（几 MB 到几十 MB）
2. 需要在各种环境下快速部署
3. 对延迟敏感但吞吐量要求不极端
4. 用户可能没有 Python 环境

想象这样一个场景：你训练了一个物理场模拟的神经网络，想让同事或客户在浏览器里交互式地探索模拟结果。传统做法是要么预渲染成视频，要么搭个服务器跑推理。Kuma 的方案是给个 HTML 文件，打开就能跑。

架构上的争议点

作者在 Reddit 帖子里抛出了几个他自己也在纠结的问题，社区讨论挺热烈。

把 kernel 嵌入 artifact 是不是坏主意？

正方观点：自包含就是优势。不依赖外部 kernel 库意味着不会因为版本不匹配出问题，部署简单。

反方观点：kernel 无法热更新。如果发现某个算子实现有 bug 或者有更优实现，得重新编译整个 artifact。另外，不同 GPU 架构可能需要不同优化策略，固定的 kernel 难以适配。

我的看法：对于科学计算这类相对稳定的场景，嵌入 kernel 利大于弊。但如果是频繁迭代的生产系统，kernel 分离可能更灵活。

和现有方案比有什么优势？

社区里有人问：为什么不直接用 ONNX.js 或者 Transformers.js？

关键区别在于编译 vs 解释。ONNX.js 是加载 ONNX 模型然后解释执行，Kuma 是提前编译成 native kernel。理论上编译方案的运行时开销更小，但目前缺乏系统性 benchmark。

另一个区别是依赖。Transformers.js 底层依赖 ONNX Runtime Web，这是个相当重的运行时。Kuma 的运行时设计上就是轻量的，因为核心计算逻辑已经编译进 artifact 了。

这到底解决的是真实问题还是想象中的问题？

这个问题很尖锐，但也很实际。

如果你的场景是大模型推理，浏览器方案目前肯定不是最优选择——显存限制、计算密度都是问题。但如果是小模型、交互式应用、离线场景，browser-native 推理是有真实需求的。

举几个例子：

• 隐私敏感应用：数据不出浏览器，推理在本地完成
• 离线工具：不依赖网络，打开就能用
• 嵌入式场景：一些 IoT 设备跑的是 Chromium 内核
• 教育演示：给学生一个 HTML 文件比配环境简单太多

实战：跑通第一个 demo

来实际操作一下。以下步骤基于 Kuma 的 GitHub 仓库。

环境准备

# 克隆仓库
git clone https://github.com/Slater-Victoroff/Kuma.git
cd Kuma

# 安装 Python 依赖
pip install -r requirements.txt

需要 PyTorch 2.x 环境。如果还在用 1.x，是时候升级了。

编译一个简单模型

假设你有一个训练好的模型：

import torch
import torch.nn as nn

class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 10)
        self.relu = nn.ReLU()
    
    def forward(self, x):
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model = SimpleNet()
model.load_state_dict(torch.load('model.pth'))

用 Kuma 编译：

from kuma import compile_model

# 导出并编译
compiled = compile_model(
    model,
    example_input=torch.randn(1, 784),
    output_path='simple_net.kuma'
)

生成的 simple_net.kuma 就是自包含的可执行包。

浏览器端运行

<!DOCTYPE html>
<html>
<head>
    <title>Kuma Demo</title>
</head>
<body>
    <script type="module">
        import { KumaRuntime } from './kuma-runtime.js';
        
        async function run() {
            const runtime = new KumaRuntime();
            await runtime.load('simple_net.kuma');
            
            // 准备输入数据
            const input = new Float32Array(784).fill(0.5);
            
            // 推理
            const output = await runtime.infer(input);
            console.log('Output:', output);
        }
        
        run();
    </script>
</body>
</html>

需要用本地服务器运行（WebGPU 不支持 file:// 协议）：

python -m http.server 8000

打开 http://localhost:8000，在控制台看输出。

调试技巧

如果跑不起来，先检查 WebGPU 支持：

if (!navigator.gpu) {
    console.error('WebGPU not supported');
}

然后检查 adapter 和 device 获取是否成功：

const adapter = await navigator.gpu.requestAdapter();
if (!adapter) {
    console.error('No GPU adapter found');
}
const device = await adapter.requestDevice();

常见问题：

• Chrome 需要开启 chrome://flags/#enable-unsafe-webgpu（某些版本）
• Safari 需要在开发者菜单里开启 WebGPU
• Firefox 的支持情况要看具体版本

性能考量

目前 Kuma 还在早期阶段，性能数据不多。但从架构上可以分析几个关键点。

首次加载

首次加载需要：

1. 下载 artifact（权重大小决定）
2. 解析计算图和元数据
3. 编译 WGSL shader（浏览器完成）
4. 分配 GPU 缓冲区

shader 编译可能是瓶颈。复杂模型有几十甚至上百个 kernel，首次编译可能需要几秒。但这是一次性开销，编译结果会被浏览器缓存。

推理延迟

推理阶段主要是 GPU 计算时间。WebGPU 的调度开销比 native CUDA 高一些，但对于中小模型来说差距不大。

和 Native 方案的对比

粗略估计，同样的模型在 WebGPU 上的性能大概是 CUDA 的 50%-80%，取决于模型类型和具体算子。卷积神经网络的差距可能小一些，Transformer 类模型的差距可能大一些（因为 attention 机制对 memory bandwidth 敏感）。

但换个角度想：WebGPU 方案不需要用户安装任何东西，打开浏览器就能跑。这个便利性在某些场景下比性能更重要。

局限性和未来方向

诚实地说，Kuma 目前还有不少限制。

算子覆盖

目前支持的算子集有限。常见的卷积、全连接、激活函数、池化应该都有，但一些特殊算子可能缺失。如果你的模型用了 custom op，大概率需要自己写 WGSL kernel。

动态形状

静态形状编译容易，动态形状难。比如处理变长序列的 RNN/Transformer，如果序列长度不固定，编译策略会复杂很多。PyTorch 2.x 的 TorchDynamo 在这方面做了很多工作，Kuma 能借鉴多少还有待观察。

内存管理

WebGPU 的内存管理和 native GPU 编程不太一样。buffer 分配、数据传输的 API 设计影响性能上限。Kuma 的运行时如何高效管理内存，是个长期优化方向。

生态整合

能否和 Hugging Face、PyTorch Hub 等模型仓库打通？能否支持常见的预训练模型？这些决定了实际使用的便利程度。

类似项目对比

值得提一下的是 jmaczan 的 torch-webgpu 项目，思路和 Kuma 有相似之处——都是 PyTorch 到 WebGPU 的桥梁。torch-webgpu 的侧重点是跨 GPU 厂商的兼容性，声称支持四个 GPU 厂商、三种后端、三种浏览器的组合。

两个项目可以互相参考。Kuma 更强调「自包含」的打包理念，torch-webgpu 更强调运行时的通用性。

写在最后

把 PyTorch 模型编译成浏览器可执行文件，这个想法本身不新鲜。但 Kuma 的方案在「自包含」这个点上做得比较彻底——一个文件包含所有依赖，扔到任何支持 WebGPU 的浏览器就能跑。

对于科学计算、教育演示、隐私敏感应用这类场景，这种部署方式确实有吸引力。但要成为主流方案，还需要解决算子覆盖、性能优化、开发者工具链等一系列问题。

项目还在早期，作者也在征集架构反馈。如果你有模型部署经验，不妨去仓库看看，说不定能贡献一些想法。

---

参考来源

• Kuma GitHub 仓库 - 项目源代码和文档
• Reddit 讨论帖 - 作者发布的原始帖子和社区讨论
• torch-webgpu 项目 - 类似思路的 PyTorch WebGPU 项目
• PyTorch 2.0 重磅发布 - PyTorch 2.x 编译器技术背景介绍