昨天聊 CUDA，今天 Google 的 TorchTPU 又提供了一个新案例 昨天提到一个问题： PyTorch 接口兼容了，原本为 CUDA 写的高性能… · @manongfu · TwKit

@manongfu

昨天聊 CUDA，今天 Google 的 TorchTPU 又提供了一个新案例 昨天提到一个问题： PyTorch 接口兼容了，原本为 CUDA 写的高性能代码怎么办？ Google 这次除了开源 TorchTPU，还推出了基于 TorchTPU 的 vLLM 和 SGLang 推理后端。 vLLM 已经有成熟的模型实现、请求调度、KV Cache 管理和 API 服务，没必要为了支持 TPU 再开发一套。 所以新后端采用独立的平台插件，尽量沿用 vLLM 原有的这些功能。 对于 PyTorch 的常规张量运算，可以通过 torch.compile 交给 XLA 编译器，生成适合 TPU 执行的代码。 但 Attention、MoE、量化矩阵乘法等对性能影响较大的计算，仅靠通用编译路径还不够，仍然需要专门优化的 TPU Kernel。 Google 在这里复用了部分基于 JAX/Pallas 编写的 Kernel，再把它们作为自定义算子接入 PyTorch。 这样一来，开发者可以继续使用 PyTorch 编写模型，底层遇到需要特殊优化的计算时，仍然可以调用 TPU 专用的实现。 从这次实现来看，为 vLLM 增加一种新芯片的支持，不一定需要把整个推理框架重新开发一遍。

曝光 8 · 评论 2 · 点赞 1 · 书签 0