Building Spyre as a Native PyTorch Device

torch.utils.rename_privateuse1_backend("spyre")
torch._register_device_module("spyre", make_spyre_module())

torch.compile
` ``

运行时管理的主机工作

与流关联的有序主机操作

本文其余部分将逐步介绍需要实际设计工作的四个映射。

## 映射 1 – 执行管道到流

运行时需要一个地方来放置有序、异步且按设备划分的工作。在 PyTorch 中，这个抽象是流。对于 Spyre 来说，有用的映射不是硬件行为完全像 PyTorch 流，而是流有序工作可以被降级为类型化的运行时操作，同时运行时保留这些操作所需的依赖关系。判断的关键在于决定边界两侧的内容。

特定于程序的工作保留在流边界之上。在 Spyre 中，面向框架的层将编译后的产物转换为准备好的启动计划：一个记录运行该产物所需的主机工作、数据移动、设备计算和同步的配方。该层可以携带张量布局和符号参数元数据，并为 PyTorch 传递的实际张量绑定启动时的值。一旦通过流提交工作，下层运行时看到的是具有显式操作数和依赖关系的类型化操作，而不是需要重新解释的程序描述。

下层不会误解程序描述，因为它从未收到过。不同的提交机制也可以位于下方，而无需更改面向框架的代码。

图 1：边界应放置的位置

保持程序知识在队列之上。所有理解编译后程序本质的内容都应属于面向框架的层。面向运行时的层接受带有显式操作数和依赖关系的类型化操作，用于主机工作、设备计算、数据移动、同步及其他运行时任务。

准备和启动发生在不同时间点。将编译产物转换为可提交操作是按产物进行的工作，而非按启动进行的工作：

准备阶段解析产物，分配并传输程序二进制文件，将执行计划转换为有序的类型化步骤列表。启动阶段遍历这些步骤并入队操作。任何在首次看到产物时能解决的内容都应在那时解决，因此启动路径仅涉及操作构造和入队。

### 为什么运行时不需要另一个图？

早期集成已通过自定义后端实现的PyTorch FX图开始。FX本身不是问题；图应位于编译器边界。不匹配发生在交接之后，当启动、复制和张量参数被转换为第二个后端特定的运行时图时。

运行一个已编译程序意味着读取并反序列化该运行时图，为其张量参数添加节点和边，加载并解析结果后再发出设备操作。主机到设备的复制同样会围绕数据转换和传输构建一个小图。这种接口在运行时需要完整图时很有用，但为那些在启动时已完全确定的操作增加了图处理工作。

将整个编译图复制到我们自己的抽象中，也导致核心运行时与PyTorch之间的通信中断。PyTorch无法得知张量是驻留在设备还是主机上。结果，PyTorch认为每个张量都是CPU张量，即使核心运行时实际上已将其移动到设备上。

新的边界将FX图保留在Inductor中，并仅对编译产物进行一次准备。之后，运行时需要的是类型化操作的有序配方，而不是每次调用时都要重建或遍历的另一个模型图。

## 映射2 – 设备内存到分配器

原生PyTorch设备不能将每个操作都视为“将输入复制到加速器，执行，再将输出复制回来”。设备为"spyre"的张量需要真实设备存储，其生命周期由PyTorch通过分配器和存储契约管理。早期的自定义torch.compile后端可以通过定制运行时路径保持特定模型状态（如权重和KV缓存条目）驻留，但这与PyTorch原生张量驻留不同。普通PyTorch值不是简单的设备为"spyre"的张量，其存储由PyTorch管理，且其分配可被下一个操作重用。这对即时执行尤其昂贵，因为每个小操作否则都会创建另一个传输边界。

Spyre 在该基础分配器协议中添加了硬件约束。框架分配器通常将设备内存作为由指针寻址的平面池呈现。在该接口之下，Spyre 以区域（region）形式管理内存——每个区域是通过句柄标识的设备内存连续块。一个区域可以包含多个张量的分配。由于句柄预算有限，尤其是在显卡在多个租户之间共享时，分配器必须将大量张量压缩到少量区域中。

Spyre 支持 PF 模式（显卡专属于单个租户）和 VF 模式（多个租户共享显卡）。VF 模式具有更严格的句柄预算。分配器会获取少量大区域，并在其中进行对齐块的子分配，使大量活跃张量共享少量句柄。

在两种模式下，每个分配都使用相同类型的描述：一个或多个区域-偏移量组合。仅区域标识符的解释方式不同，因此分配器之上各层无需根据部署模式进行分支判断。分配器返回该描述及释放方式；PyTorch 的引用计数机制会在张量活跃期间保持其在设备上的存在。

完整的分配器接口增加了 PyTorch 标准内存管理 API，包括使用统计、缓存控制和流关联。该接口很有用，但并不是使张量驻留的关键。

此处的分配描述包含的不仅是地址：基本可寻址单元由区域标识符、该区域内的偏移量和长度组成。跨内存域交错的张量包含多个此类单元，而非单一平面指针。PyTorch 为此预留了空间：分配可携带一个不透明上下文与指针一同存在，描述信息存储于此，引用计数仍负责清理工作。

图2：设备分配的真实本质

分配是一个描述信息，而不仅是地址。绑定到单一内存域的张量是一个区域-偏移量单元；跨域交错的张量包含多个此类单元。PyTorch 将该描述信息存储在不透明上下文中，其引用计数机制会通知后端分配器何时可以释放存储。

相同的表示方式也可覆盖非统一内存架构的硬件。在该类设备中，内存被划分为多个域，核心访问不同域的效率存在差异。张量可绑定到使用它的核心附近，或跨域交错以利用总带宽。这是单设备内的局部性，而非多卡分布。

由于编译代码可能依赖于内存布局，因此内存布局必须显式且具有约束力。因此，所有携带分配信息的层级也必须携带该分配是绑定到单一域还是跨域分布的信息。

驻留机制也使内存布局成为显式契约。模型适配器会将权重和KV缓存放置在编译内核预期的布局中，而编译器会插入中间张量的合法布局转换。这些工作并未消失，而是从隐藏的运行时图转移到了模型准备、分配元数据和编译器中。

设备放置与卸载是两个独立的问题。放置决定了张量在加速器内部的存储位置。卸载则将暂时不需要的状态（如KV缓存页）转移到次级存储层级，并在使用前恢复这些状态。该层级可能是主机内存，也可能是服务系统层级的SSD或网络附加存储。运行时接口应描述数据传输过程，而无需暴露该层级背后的物理介质。

## 映射3 – 预编译程序到启动时的工作

每个预编译加速器都必须在某处解决这个问题。编译器了解计算结构，但不知道调用方数据的位置，这两者必须在程序运行前达成一致。随着运行时的发展，Spyre采用了三种方法。当前运行时已完全采用第三种方法——运行时修补。

早期的图中介方法让编译任务拥有地址。输入数据被复制到任务拥有的设备缓冲区，结果也被复制回主机。这些缓冲区可以在多次启动间保持具体地址，因为它们的生命周期和位置由任务控制。这种方法仍适用于任务完全拥有的内存（如中间结果和编译器管理的临时缓冲区），但无法推广到由PyTorch分配器选择存储位置的驻留张量。

中间运行时将程序的地址窗口指向张量。设备翻译（xlat）表暴露了少量地址窗口，因此运行时将每个窗口绑定到一个分配器放置的张量。无需修补，张量也不需要复制到编译器管理的位置。限制因素是输入参数数量而非内存分配：每个窗口对应一个张量，且窗口数量有限，因此大型融合内核无法接收所有张量。

当前运行时在启动时修补程序。地址在分配器放置张量前保持符号化。这避免了重复的输入输出复制，也突破了窗口数量上限，成为处理分配器放置张量的通用路径。在Spyre中，编译器生成一个布局固定但张量起始地址可能为符号化的工件；二进制文件在地址位置保留占位变量。替换真实值的过程包含三个普通步骤，与其他流程完全一致：

- 主机回调在CPU上运行，将解析后的地址写入固定主机缓冲区。
- 传输将该缓冲区复制到程序自己的内存分配区域，偏移量由编译器指定。
- 计算运行时读取这些值并修补其操作数。

在单个流中，完成顺序足以防止三个步骤发生竞争。在接下来描述的初始双流设计中，主机回调和传输在准备流中保持顺序。计算操作在设备流上运行，并等待传输完成后记录的事件。这在保持相同顺序的同时，允许准备阶段与更早的计算操作重叠。

队列排序控制操作的发出时机，但不控制在它们之前运行的硬件。Spyre的程序分发机制正是如此：向指令缓冲区提供数据的单元会积极获取数据，而不是等待当前程序执行完毕，因此可以在补丁写入内存之前，将即将被修补的程序拉入缓冲区并运行过时的字节。解决方案是在程序分发机制下一层设置屏障，将获取操作延迟到补丁落地后。队列排序控制提交顺序；下层的屏障防止了读取在写入之前的竞争。

图3：启动时修补的三个有序步骤

启动时修补包含三个普通步骤和两个保证。主机回调将解析后的地址写入固定缓冲区，传输操作将它们移动到程序分配的编译器指定偏移位置，设备读取这些地址并修补操作数。准备流对主机回调和传输操作进行排序，设备流在传输完成前等待计算。程序分发中的屏障也防止获取引擎在补丁落地前读取程序。

初始的双流设计可以复用一个固定主机中转缓冲区，因为对该缓冲区的所有写入和传输操作都保留在同一准备流中。完成排序可防止下一个主机回调在前一个传输操作消耗缓冲区之前覆盖主机缓冲区。这一保证无法自动扩展到多个准备流或多级主机预取。这些情况需要为每个并发传输分配单独的主机缓冲区，或使用完成门控池，直到传输完成才允许复用。

设备修正区域是独立的。传输操作写入该区域，计算操作读取该区域，因此事件在流之间对这些操作进行排序。如果下一次启动复用相同的设备区域，其传输操作也必须等待前一次计算完成读取。主机准备仍可与该计算操作重叠；传输操作仅在目标与计算使用的内存无关时才能与计算操作重叠。

启动时修补是当前通用解决方案。其主机准备可通过下文描述的多流机制与设备计算重叠，传输操作在目标为独立设备存储时也可重叠。编译器拥有的地址和地址窗口绑定可能仍适用于特定场景，但运行时不会为普通张量输入在三种方案中进行选择。

未来支持将允许运行时通过与地址修正相同的启动时修正机制，提供可变张量维度的实际大小。当使用符号化维度编译时，一个工件可以在其支持的范围和布局约束内服务于多个输入大小，而无需为每个大小重新编译。

## 映射4 – 独立流水线到并发流

传输和计算使用独立的硬件流水线，但单个流仍可能让其中一个等待另一个。在单流序列中，主机准备修正数据，传输将其发送到设备，计算在下一次迭代开始前消耗这些数据。这是正确的，但也会让主机和传输流水线在设备计算时处于等待状态。

在一次迭代中，主机回调必须在数据传输读取校正缓冲区之前完成，而数据传输必须在计算使用校正后的程序之前完成。这可以防止数据传输读取不完整的缓冲区，以及计算读取未打补丁的程序。在不同迭代之间，迭代N+1的准备工作可以在迭代N进行计算时开始，因为其地址在启动时已经确定。

运行时会保持必要的顺序，并将独立的工作转移到其他流中。准备流可以在设备流计算迭代N的同时，为迭代N+1生成校正数据。事件在实际依赖点连接流：计算N+1等待其数据传输完成，但与计算N无关的主机准备工作不需要等待计算N。当数据传输目标是独立的设备存储时，数据传输也可以与计算重叠。重复使用相同的设备校正区域会引入反向依赖，因此特定的数据传输需要等待之前的计算完成读取。

事件通过PyTorch读者已知的原始方法连接队列：一个流记录标记，另一个流等待该标记。准备流发出信号表明数据传输已完成；设备流在计算前等待该信号。在关键位置，顺序保持完整，通过队列之间的表达而非单个队列内部实现。

在运行时层，事件和同步操作与CUDA紧密对应：

CUDA

此处

排队一个事件信号操作；只有当该流上所有已排队的操作都完成时，信号才会触发

排队一个事件等待操作；队列会在该等待点暂停，直到事件被触发，如果生成队列先终止则直接失败

同步队列：阻塞直到队列清空，并重新抛出该队列记录的第一个错误

在CUDA中，调用者显式记录和等待。此处运行时可以为那些设备内存读写足迹已知的操作推导出边。框架层将准备和设备工作路由到适当队列；当这些设备范围冲突时，运行时会插入信号和等待操作。程序和张量分配参与该模型。主机暂存缓冲区复用是独立的生命周期规则，由同一流完成顺序或完成门控池保护。较低层运行时支持显式信号和等待操作，用于无法从设备内存足迹推断的依赖。这些操作目前是内部实现；通过PyTorch的用户可调用事件记录和等待功能已计划但尚未实现。目前，用户可以通过stream.synchronize()或torch.accelerator.synchronize()等待完成，这会阻塞主机线程。

## 映射带来的好处与成本

### 用户体验

用户接触的界面是他们已经熟悉的。`tensor.to("spyre")`、`torch.compile`和`torch.Stream`使用与其他设备相同的PyTorch概念。这并不自动提供操作符覆盖、性能分析、分布式集合操作、CUDA特定扩展或模型特定内核质量；这些工作仍属于后端实现。

对于由编译支持的即时模式操作，即时模式并非第二个实现。Torch-spyre围绕编译期间使用的相同分解方法注册了一个PrivateUse1内核。在`torch.compile`的追踪过程中，包装器直接调用该分解方法，使PyTorch能够将其捕获到周围的FX图中。在追踪之外，第一次即时调用会编译并缓存该分解方法；后续调用将复用已编译的入口点。因此，这些操作复用相同的Inductor和启动机制，而无需维护额外的手写运行时。

该选择仍然暴露了首次使用的编译过程；不支持的操作可能失败或回退，不同的张量范围可能需要另一个编译产物。其优势在于，修复和回归的修改适用于两种模式的单一执行路径。

由编译支持的即时操作也是可能的最小编译程序：一个操作、真实张量、真实设备内存和真实启动。如果模型失败但对应的即时操作通过，问题更可能出在组合而非操作本身。单个操作的测试可以验证整个编译和启动堆栈，因此每个操作的测试套件同时充当集成测试套件。

由于驻留性是通用的而非针对每种情况手写，张量无需为每种类型的张量编写专门的驻留代码即可保留在设备上。这正是即时模式实用的关键所在。

### 性能

性能论证的核心在于哪些操作被移出了启动路径。将类型化操作提交到队列，而非构建图来描述它们，可以消除每个启动过程中的一次磁盘读取、一次反序列化、一次按参数缝合以及一次编译和解析。启动过程因此简化为构建几个操作并将其加入队列。

在真实模型中，这种优化效果会叠加，因为这些成本原本是每个启动都要支付的。在Granite 3.3 8B模型上，当批量大小为1且序列长度为1024时，仅将图构建移出传输路径就能带来9.9%的预填充加速和13.6%的解码加速——这仅仅是改变了字节移动方式的改动。若进一步将图构建移出计算调度路径，预填充和解码速度分别提升至原来的1.7倍和2.4倍，其中约四分之三的性能提升来自计算路径。

如第4部分映射中所述的并发性是另一个机会：将所有主机准备操作移出计算路径，可显著减少设备空闲时间。

### PyTorch本身未提供的功能

PyTorch提供了设备事件抽象，但上述跨流排序必须在操作实际执行的位置（即运行时调度器）强制实施——而运行时调度器是能看到这些操作设备内存占用的层级。因此，运行时首先构建软件事件原语和派生依赖关系。通过PyTorch的公共事件钩子暴露等效的用户创建事件，需要额外的集成工作。

### 我们回馈的内容

基于框架的扩展点进行开发并非单向交易。PyTorch核心开发者维护着openreg，这是上游参考后端和加速器模拟器，展示了无树设备如何通过PrivateUse1注册接口而无需修改核心PyTorch，该项目得到了硬件生态系统合作伙伴的贡献。在该仓库提交改进，使像torch-spyre这样的后端工作能为后续设备团队缩短开发路径；参与PyTorch技术咨询委员会，则能让这类设备的约束条件传递给下一个扩展点的设计者。我们与PyTorch的跨仓库CI中继（CRCR）集成，将上游变更在Spyre上进行测试，并将结果报告给PyTorch的CI仪表板，帮助在框架演进过程中捕获兼容性回归。我们还构建了可复用的工具，用于选择和适配上游测试，计划将其提交上游，让其他无树加速器也能受益。一个只消费接口的后端，不会让接口变得比它原本更好。

## 结论

The win is that neither Spyre nor PyTorch has to pretend to be something else. Spyre keeps its dataflow execution model, while users get a native PyTorch device with familiar tensor placement, streams, and compilation.

对于 Spyre 和 PyTorch 来说，最大的优势在于无需互相伪装成其他形态。Spyre 保留其数据流执行模型，而用户则可以获得原生的 PyTorch 设备，拥有熟悉的张量放置、流和编译方式。

For Spyre, PyTorch’s interfaces provide generic tensor residency, one compile-and-launch path for eager and compiled execution, lower launch overhead, and a clear way to overlap host preparation with device work. They also leave less backend-specific machinery to maintain. Dispatch, storage lifetime, and stream management follow framework contracts, making the runtime easier to evolve and less likely to develop inconsistent execution paths.

对 Spyre 来说，PyTorch 的接口提供了通用的张量驻留能力，为即时执行和编译执行提供统一的编译-启动路径，降低启动开销，并清晰地实现了主机准备与设备工作之间的重叠执行。此外，它们减少了需要维护的后端特定机制。调度、存储生命周期和流管理遵循框架契约，使运行时更容易演进，也更不容易产生不一致的执行路径。

For PyTorch, Spyre demonstrates that those extension points can support hardware that is not organized around a conventional GPU execution model. The lessons that flow back through openreg and the PyTorch Technical Advisory Council make those interfaces more useful to the next accelerator backend as well. The result is a more robust Spyre runtime and a PyTorch ecosystem that can accommodate a broader range of hardware without giving each device a separate framework.

对 PyTorch 来说，Spyre 展示了这些扩展点可以支持不基于传统 GPU 执行模型的硬件。通过 openreg 和 PyTorch 技术咨询委员会反馈的经验，这些接口对下一个加速器后端也变得更加有用。最终结果是更健壮的 Spyre 运行时，以及一个无需为每种设备单独设计框架即可适应更广泛硬件的 PyTorch 生态系统。

/post-content

/inner-wrap