首页
学习
活动
专区
圈层
工具
发布

IBM Spyre加速器的PyTorch原生支持构建之路

我们最近发布了在PyTorch生态系统中启用IBM Spyre加速器的2026年上半年路线图。在配套的技术说明中,我们介绍了该硬件的特性——32个活跃AI核心、混合精度SIMD-脉动阵列以及可编程数据流设计。

在这里,我们想详细介绍路线图的关键主题,以及我们如何在一个主要由GPU塑造的生态系统中为数据流加速器构建一流的PyTorch支持。我们的理念是生态优先——我们利用上游机制,最小化自定义代码,并回馈那些让下一个加速器更容易遵循相同路径的组件。

扩展Inductor以支持数据流抽象

我们正在树外扩展inductor,以处理数据流加速器所需的抽象。三个扩展最为重要:首先,我们引入了基于块的张量布局,使编译器能够推理Spyre环形连接核心所期望的块结构数据移动。其次,我们添加了多核工作划分过程,在编译期间而非运行时将块分配到Spyre的32个核心上。第三,我们添加了暂存器优化——Spyre核心使用显式管理的片上内存而非硬件缓存,inductor在调度数据时需要考虑这一点。这些扩展将使torch.compile能够为我们2026年上半年范围内的每个优先模型生成高效的Spyre代码,从Llama 3.1 8B到Granite 4 Hybrid 30B。

后端编译器中间表示的两阶段方法

我们对位于inductor高级图和Spyre机器码之间的后端编译器中间表示采取两阶段方法。

在第一阶段,SuperDSC(SDSC)将作为后端编译器中间表示——所有操作降级和代码生成的单一入口点。我们优先模型所需的每个torch操作都可以在SDSC中表达,在PyTorch集成层和硬件特定优化之间提供清晰的分离。

在第二阶段,我们将过渡到KernelTile IR(KTIR),这是一个更符合社区的规范,与TileIR等新兴计划更加一致。KTIR将泛化块级表示,使其他数据流加速器——不仅仅是Spyre——可以将其用于更低级别的调度和代码生成。我们计划在今年上半年发布完整的KTIR规范,并且我们正在设计的开源调度算法可以适应我们自己硬件之外的应用。

设备集成与分布式通信

Spyre将完全通过树外扩展注册为PyTorch设备:设备生命周期、内存管理、数据传输和调度。我们的目标是以这种方式处理100%的注册,与直接设备访问相比开销低于5%。我们计划将构建的通用原语回馈到PyTorch核心的OpenReg测试基础设施中。

对于多卡推理,我们通过torch.inductor编译功能性集合操作(all-reduce、all-gather),这将使我们在2026年上半年实现所有优先模型的分布式推理。长期来看,我们计划迁移到torch.distributed,并最终随着社区通信层的稳定迁移到torch.comms。

生产推理与vLLM集成

生产推理将通过vLLM运行。我们将Spyre作为vLLM平台插件启用,采用上游模型实现而非维护我们自己的分支。我们的优先模型将通过Spyre上的vLLM进行端到端服务。

两项优化将推动实际可用性。新的Spyre注意力后端将消除同质序列长度约束,直接减少Token间延迟。上游vLLM中改进的torch.compile工件缓存将把启动时间降至几秒钟。我们正在与vLLM社区合作,以稳定平台插件接口。

测试与持续集成

我们正在构建一个分层测试金字塔,将验证整个堆栈:操作级正确性、inductor编译和降级、模块级测试(包括注意力、归一化和激活)、顶层模型质量和性能,以及端到端vLLM推理。所有测试都将限定在优先模型范围内,并每晚运行,回归失败将在数小时内标记。

我们正在将此持续集成基础设施构建为对PyTorch生态系统的树外贡献,建立其他加速器团队可以采用的模式。我们的目标是每晚运行的通过率超过95%,完整流程在三小时内完成。

回馈社区

生态优先意味着回馈,而不仅仅是在上面构建。本半年有三项突出贡献:我们计划上游OpenReg原语,使树外设备测试成为PyTorch的一流能力。我们正在努力将KTIR泛化为社区规范,使数据流加速器共享通用的块级中间表示,而不是各自发明自己的。我们将记录树外持续集成模式,使下一个加速器团队不必从头解决基础设施问题。

设计文档和RFC存放在我们的公共存储库中。我们欢迎参与——无论是审查KTIR规范、试用Spyre跟踪分析器,还是参与关于PyTorch原生加速器支持应该是什么样子的讨论。

Q&A

Q1:IBM Spyre加速器有什么硬件特点?

A:IBM Spyre加速器拥有32个活跃AI核心、混合精度SIMD-脉动阵列以及可编程数据流设计。其核心使用显式管理的片上内存而非硬件缓存,核心之间采用环形连接架构,适合块结构数据移动。

Q2:PyTorch如何支持Spyre这样的数据流加速器?

A:通过树外扩展inductor来处理数据流加速器所需的抽象,包括引入基于块的张量布局、添加多核工作划分过程以及暂存器优化。Spyre将完全通过树外扩展注册为PyTorch设备,处理设备生命周期、内存管理、数据传输和调度,目标是与直接设备访问相比开销低于5%。

Q3:KTIR在Spyre的PyTorch支持中扮演什么角色?

A:KTIR(KernelTile IR)是第二阶段的后端编译器中间表示,它是一个更符合社区的规范,与TileIR等新兴计划一致。KTIR将泛化块级表示,使其他数据流加速器也可以将其用于更低级别的调度和代码生成,而不仅仅服务于Spyre硬件。

  • 发表于:
  • 原文链接https://page.om.qq.com/page/ObGljQj3f_wgZS85eeEN3OaA0
  • 腾讯「腾讯云开发者社区」是腾讯内容开放平台帐号(企鹅号)传播渠道之一,根据《腾讯内容开放平台服务协议》转载发布内容。
  • 如有侵权,请联系 cloudcommunity@tencent.com 删除。
领券