IBM 如何将 Spyre 打造为原生 PyTorch 设备

TL;DR

Spyre becomes a native PyTorch device by connecting PyTorch’s existing device, allocator, stream, and compiler abstractions through torch-spyre to the Spyre runtime and firmware. PrivateUse1 gives Spyre a real device identity, PyTorch allocator storage keeps tensors resident on device="spyre", streams provide familiar ordering semantics, and runtime events connect dependent work across streams. FX graphs stay in the Inductor compiler path. This post focuses on the PyTorch-to-Spyre device/runtime integration: how device identity, allocator-backed storage, streams, events, and compiled-program launch map onto Spyre’s runtime constraints. Once a compiled artifact exists, launch is a prepared recipe of typed operations over ordered queues rather than a graph-based runtime interface. The result is one path for eager and compiled execution, lower launch overhead, and a PyTorch surface that maps cleanly onto Spyre’s hardware.

Spyre is IBM’s dataflow AI accelerator, optimized for inference. It is designed for enterprise teams running AI alongside their applications and data on IBM Z, LinuxONE, and Power systems. Its reduced-precision compute is well suited to the matrix-heavy work in language generation and embedding models.

For developers, the challenge is using that hardware through familiar PyTorch code while keeping up with new models and inference techniques. Efficient execution needs more than a compiler: tensors must be able to stay on the device between operations, launches need to be lightweight, and host work and data transfers should overlap computation wherever dependencies allow. This post explains how IBM’s torch-spyre team builds on the PyTorch community’s interfaces to support those workflows on a dataflow accelerator.

The dataflow constraints behind the PyTorch surface

The card has 32 cores connected by a high-bandwidth ring. Each core has 2 MB of local scratchpad and arrays of processing elements. Up to 128 GB of LPDDR5 provides storage for tensors and programs.

Those two levels of memory have different owners. The runtime stack allocates and manages LPDDR5. Before a program runs, the compiler has already emitted the loads and stores that move tiles between LPDDR5 and each core’s scratchpad. Data moves across that path in 128-byte sticks, so layout and alignment affect how efficiently it can be fetched. The runtime does not allocate or schedule the scratchpad directly.

A native PyTorch device still has to honor the execution model of the hardware underneath it. For Spyre, that means mapping PyTorch’s device, allocator, stream, event, and launch abstractions onto a runtime built around compiled programs, ordered queues, explicit data movement, and fixed layout contracts. Four constraints drive the design.

Computation is triggered by data, not by a single kernel instruction stream. A GPU kernel is typically organized around scheduled groups of threads executing the kernel’s instruction sequence. A compiled Spyre kernel is arranged differently: it contains programs for multiple functional units – such as processing elements, special-function pipelines, and load/store units – and each unit runs the program compiled for that unit. Hardware queues and local memories connect those units, so progress is driven by operands becoming available at the consuming unit. The Spyre compiler builds that device-side producer/consumer schedule; the runtime submits the compiled kernel as a complete device-compute operation with its tensor arguments.

Runtime overlap comes from separate pipelines, not concurrent compute launches. Spyre has a compute pipeline and data-movement pipelines, so a transfer can run while a compiled program is computing. Those pipelines are not interchangeable: a compiled program can use many cores in parallel, but independently submitted compute operations share one runtime compute queue and wait behind each other. Runtime-level overlap therefore comes from moving data while computing, not from launching two compute operations at once.

Ordered queues preserve per-stream completion order. Work arrives as a sequence of typed operations – move this in, run this, move that out – and each stream completes those operations in order. To use the separate hardware pipelines without weakening that guarantee, the runtime assigns independent transfer and compute work to different streams and connects them with an event where one consumes the other’s result. Putting all work on one stream preserves correctness but serializes those pipelines.

Programs are compiled ahead of time against a fixed layout. The compiler fixes how values compose into the device’s native fixed-size chunks and what order dimensions appear in, and it can fix addresses for memory belonging to the job itself – intermediates and compiler-owned temporary buffers. What it cannot fix is anything whose address the caller owns: tensor inputs and outputs stay symbolic because the allocator has not placed them yet. The runtime must supply those addresses at launch.

In the implementation described here, PyTorch does not talk to the card directly. torch-spyre provides the PyTorch device integration, while the Spyre runtime and firmware manage device work and communication. That layering is why the rest of this post separates PyTorch concepts such as tensors, streams, and events from lower-level concerns such as device regions, DMA-visible host staging buffers, and explicit work submission.

The PyTorch abstractions that Spyre maps onto

PyTorch already has concepts that map onto those constraints.

A device identity. PrivateUse1 registration starts with two calls:

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

Those calls give PrivateUse1 the name spyre and register its device module. With the allocator, copy kernels, and dispatcher registrations behind that identity, tensor.to("spyre") works and the dispatcher routes Spyre tensors to Spyre kernels.

Device-management APIs. Backend-specific controls live under torch.spyre, while torch.accelerator gives portable code a common interface for availability, device selection, streams, synchronization, and memory reporting. Tensor allocation and kernel launch still flow through tensors, the dispatcher, the device allocator, and the backend runtime.

A stream abstraction. Implementing the device-guard hooks the framework asks for is what makes native torch.Stream(device="spyre") and torch.stream(...) work on your device. The framework-native surface works because those hooks are there, not because the backend reimplemented PyTorch’s stream API.

A device allocator interface. Its core contract hands back device storage and a way to release it. Additional hooks connect the backend to PyTorch’s memory statistics, cache controls, and stream association.

Reference-counted tensor storage. Tensor lifetime becomes the framework’s problem: when the last reference to a Spyre tensor goes away, PyTorch calls back into the backend allocator to release or recycle the device allocation.

A compiler path through Inductor. torch.compile(backend="inductor") keeps FX graphs in PyTorch’s compiler pipeline. The same lowering, cache, and launch path can serve compiled models and registered eager operations.

An event model. PyTorch defines events as a way to express ordering between streams. Spyre uses the same model internally through runtime events and derived dependency edges.

Laid side by side, the correspondence is close:

The rest of this post walks through the four mappings that take real design work.

Mapping 1 – execution pipelines to streams

The runtime needs somewhere to put work that is ordered, asynchronous, and per-device. In PyTorch, that abstraction is a stream. For Spyre, the useful mapping is not that the hardware behaves exactly like a PyTorch stream; it is that stream-ordered work can be lowered into typed runtime operations while the runtime preserves the dependencies those operations require. What takes judgment is deciding what sits on either side of that boundary.

Program-specific work stays above the stream boundary. On Spyre, the framework-facing layer turns the compiled artifact into a prepared launch plan: a recipe that records the host work, data movement, device compute, and synchronization needed to run that artifact. That layer can carry tensor layout and symbolic argument metadata and bind launch-time values for the actual tensors passed by PyTorch. Once work is submitted through a stream, the lower runtime sees typed operations with explicit operands and dependencies, not a program description to reinterpret.

The lower layer cannot misinterpret a program description because it never receives one. A different submission mechanism can also sit underneath without changing the framework-facing code.

Figure 1: Where to put the boundary

Keep the program knowledge above the queue. Everything that understands what a compiled program is belongs in the framework-facing layer. The runtime-facing layer accepts typed operations with explicit operands and dependencies for host work, device compute, data movement, synchronization, and other runtime tasks.

Preparation and launch also happen at different times. Translating a compiled artifact into submittable operations is per-artifact work, not per-launch work:

# once, when the compiled artifact is first seen
job_plan = prepare_kernel(spyrecode_dir)

# per invocation
launch_jobplan(job_plan, args)

Preparation parses the artifact, allocates and transfers the program binary, and translates the execution plan into an ordered list of typed steps. Launch walks those steps and enqueues operations. Anything that can be resolved when the artifact is first seen should be resolved there, so the launch path is only operation construction and enqueue.

Why not another graph at runtime?

The earlier integration already started with a PyTorch FX graph through a custom backend implementation. FX was not the problem; graphs belong at the compiler boundary. The mismatch came after that handoff, when launches, copies, and tensor arguments were translated into a second backend-specific runtime graph.

Running one already-compiled program meant reading and deserializing that runtime graph, adding nodes and edges for its tensor arguments, and loading and parsing the result before issuing device operations. A host-to-device copy similarly built a small graph around data conversion and transfer. That interface was useful when the runtime needed a whole graph, but it added graph-processing work to operations that were already fully determined by launch time.

Copying the entire compiled graph into our own abstractions also caused a break in communication between the core runtime and PyTorch. PyTorch had no visibility as to whether tensors were resident to the device or to the host. As a consequence, PyTorch believed every tensor was a CPU tensor, even if the core runtime had really moved it to the device.

The new boundary keeps FX graphs in Inductor and prepares the compiled artifact once. After that, the runtime needs an ordered recipe of typed operations, not another model graph to reconstruct or traverse on every invocation.

Mapping 2 – device memory to the allocator

A PyTorch-native device cannot treat every operation as “copy inputs to the accelerator, run, and copy outputs back.” A tensor on device="spyre" needs real device storage whose lifetime PyTorch owns through its allocator and storage contracts. The earlier custom torch.compile backend could keep specific model state such as weights and KV-cache entries resident through bespoke runtime paths, but that was not the same as PyTorch-native tensor residency. An ordinary PyTorch value was not simply a device="spyre" tensor whose storage PyTorch managed and whose allocation could be reused by the next operation. That is especially costly for eager execution, where each small operation would otherwise create another transfer boundary.

Spyre adds a hardware constraint to that basic allocator contract. Framework allocators commonly present device memory as a flat pool addressed by pointers. Underneath that interface, Spyre manages memory in regions – each a contiguous chunk of device memory identified by a handle. A region can contain allocations for many tensors. Because the handle budget is limited, especially when a card is shared between tenants, the allocator must fit many tensors into a small number of regions.

Spyre supports PF mode, where a card is dedicated to one tenant, and VF mode, where several tenants share it. VF mode has the tighter handle budget. The allocator acquires a small number of large regions and sub-allocates aligned blocks within them, allowing many live tensors to share a handful of handles.

Every allocation uses the same kind of description in either mode: one or more region-and-offset pieces. Only the interpretation of the region identifier changes, so the layers above the allocator do not branch on the deployment mode. The allocator returns that description and a way to free it; PyTorch’s reference counting keeps the tensor on the device while it is live.

The fuller allocator interface adds PyTorch’s standard memory-management API, including usage statistics, cache controls, and stream association. It is useful, but it is not what makes tensors resident.

An allocation here is described by more than an address: the basic addressable piece is a region identifier, an offset within that region, and a length. A tensor interleaved across memory domains is several such pieces, not one flat pointer, and PyTorch has room for that: an allocation can carry an opaque context alongside its pointer, so the description rides there and reference counting still does the cleanup.

Figure 2: What a device allocation really is

An allocation is a description, not just an address. A tensor bound to one memory domain is a single region-and-offset piece; a tensor interleaved across domains is several. PyTorch carries that description in an opaque context, and its reference counting tells the backend allocator when the storage can be released.

The same representation can cover hardware with non-uniform memory. On such a device, memory is divided into domains and a core reaches some domains more efficiently than others. A tensor can be bound near the cores that use it or interleaved across domains to use their combined bandwidth. This is locality within one device, not multi-card distribution.

Placement must be explicit and binding because compiled code may depend on it. Every layer that carries an allocation therefore also needs to carry whether the allocation is bound to one domain or spread across several.

Residency also makes layout an explicit contract. Model adapters place weights and KV caches in layouts compiled kernels expect, while compiler passes insert legal layout conversions for intermediate tensors. The work did not disappear; it moved from a hidden runtime graph into model preparation, allocation metadata, and the compiler.

Device placement and offload are separate concerns. Placement determines where a tensor lives within the accelerator. Offload moves state that is not immediately needed, such as KV-cache pages, to a secondary storage tier and restores it before use. That tier may be host memory or, at the serving-system level, SSD or network-attached storage. The runtime interface should describe the transfer without exposing the physical medium behind that tier.

Mapping 3 – ahead-of-time programs to launch-time work

Every ahead-of-time accelerator has to solve this somewhere. The compiler knows the structure of the computation but not where the caller’s data will be, and the two have to be reconciled before the program runs. Spyre has used three approaches as the runtime has evolved. The runtime now uses the third approach, runtime patching, exclusively.

The earlier graph-mediated path let the compiled job own the addresses. Inputs were copied into device buffers owned by the job, and results were copied back out. Those buffers could keep concrete addresses across launches because their lifetime and placement belonged to the job. This still suits memory the job owns outright, such as intermediates and compiler-owned temporary buffers, but it does not generalize to resident tensors whose storage is chosen by PyTorch’s allocator.

An interim runtime pointed the program’s address windows at tensors. The device translation, or xlat, table exposes a small set of address windows, so the runtime bound each window to one allocator-placed tensor. Nothing was patched and the tensors did not need to be copied into compiler-owned locations. The limit was input arity rather than allocation: with one tensor per window and only a handful of windows, large fused kernels could not accept all of their tensors.

The current runtime patches the program at launch. Addresses remain symbolic until the allocator has placed the tensors. This avoids both repeated input and output copies and the window-count ceiling, making it the generic path for allocator-placed tensors. On Spyre the compiler emits an artifact whose layout is fixed but whose tensor start addresses may be symbolic; the binary carries placeholder variables where addresses will go. Substituting real values is three ordinary steps in the same plan as everything else:

1. A host callback runs on the CPU and writes the resolved addresses into a pinned host buffer.
2. A transfer copies that buffer into the program’s own allocation, at an offset the compiler specified.
3. The computation runs, reads those values, and patches its operands.

On one stream, completion ordering is enough to keep the three steps from racing. In the initial two-stream design described next, the host callback and transfer remain ordered on one preparation stream. The compute operation runs on the device stream and waits for an event recorded after the transfer completes. This preserves the same ordering while allowing preparation to overlap with an earlier compute operation.

Queue ordering controls when operations are issued. It does not control hardware that runs ahead of them. Spyre’s program distribution does exactly that: the unit feeding instruction buffers fetches eagerly rather than waiting for the current program to finish, so it can pull the program that is about to be patched into a buffer before the patch reaches memory, and then run stale bytes. The remedy is a barrier one level down, in program distribution, holding the fetch until the patch has landed. The queue orders submission; the lower-level barrier prevents the read-before-write race.

Figure 3: Launch-time patching as three ordered steps

Launch-time patching is three ordinary steps and two guarantees. A host callback writes resolved addresses into a pinned buffer, a transfer moves them into the program’s allocation at a compiler-specified offset, and the device reads them and patches its operands. The preparation stream orders the host callback and transfer, and the device stream waits for the transfer before computing. A barrier in program distribution also keeps the fetch engine from reading the program before the patch lands.

The initial two-stream design can reuse one pinned host staging buffer because every write to that buffer and every transfer from it stays on the same preparation stream. Completion ordering prevents the next host callback from overwriting the host buffer until the previous transfer has consumed it. This guarantee does not extend automatically to multiple preparation streams or deeper host run-ahead. Those cases need a separate host buffer for each concurrent transfer or a completion-gated pool that prevents reuse until the transfer finishes.

The device correction area is separate. The transfer writes it and compute reads it, so an event orders those operations across streams. If the next launch reuses that same device area, its transfer must also wait until the previous compute has finished reading it. Host preparation can still overlap that compute; a transfer overlaps compute only when its destination is independent of the memory that computation is using.

Launch-time patching is the current general solution. Its host preparation can overlap device compute using the multiple streams described next, and transfers can overlap when they target independent device storage. Compiler-owned addresses and address-window binding may still fit specialized cases, but the runtime does not choose among all three for ordinary tensor inputs.

Future support would allow the runtime to supply the actual sizes of variable tensor dimensions through the same launch-time correction mechanism used for addresses. When compiled with those dimensions left symbolic, one artifact could serve multiple input sizes within its supported bounds and layout constraints, without recompiling for each size.

Mapping 4 – independent pipelines to concurrent streams

Transfer and compute use separate hardware pipelines, but a single stream can still leave one waiting on the other. In the single-stream sequence, the host prepares correction data, the transfer moves it to the device, and compute consumes it before the next iteration begins. This is correct, but it also leaves the host and transfer pipeline waiting while the device computes.

Within an iteration, the host callback must finish before the transfer reads the correction buffer, and the transfer must finish before compute uses the corrected program. That prevents the transfer from reading an incomplete buffer and compute from reading an unpatched program. Across iterations, preparation for iteration N+1 can begin while iteration N computes because its addresses are already known at launch.

The runtime keeps the required ordering and moves independent work onto another stream. The preparation stream can produce correction data for iteration N+1 while the device stream computes iteration N. An event joins the streams at the actual dependency: compute N+1 waits for its transfer, but unrelated host preparation does not wait for compute N. A transfer can also overlap compute when it targets independent device storage. Reusing the same device correction area adds a reverse dependency, so that particular transfer waits for the earlier computation to finish reading it.

Events connect the queues using a primitive a PyTorch reader already knows: a marker one stream records and another waits on. The preparation stream signals that a transfer is complete; the device stream waits for that signal before computing. Ordering stays total where it matters, expressed between queues rather than within one.

At the runtime layer, event and synchronization operations map closely to CUDA:

In CUDA the caller records and waits explicitly. Here the runtime can derive an edge for operations whose device-memory read and write footprints are available. The framework layer routes preparation and device work to the appropriate queues; the runtime inserts signal and wait operations when those device ranges conflict. Program and tensor allocations participate in that model. Host staging-buffer reuse is a separate lifetime rule, protected by same-stream completion ordering or a completion-gated pool. The lower runtime supports explicit signal and wait operations for dependencies it cannot infer from device-memory footprints. These operations are currently internal; user-callable event recording and waiting through PyTorch are planned but not yet implemented. For now, users can wait for completion with stream.synchronize() or torch.accelerator.synchronize(), which block the host thread.

The initial arrangement uses two queues in a producer/consumer relationship: a preparation queue carrying host work and its transfers, a device queue carrying compute. Nothing about the mechanism stops at two, and nothing forces the split either – a launch with no host work to do never needs it. The same primitives extend to any number of preparation queues, so host work with no reason to be serial can proceed in parallel; and because transfers and compute are separate pipelines, a transfer targeting independent device storage can be in flight while the current iteration computes. Additional preparation queues must use independent staging buffers or a completion-gated pool so one queue cannot overwrite a buffer another transfer is still reading.

Host preparation becomes more visible as launch-time work grows and device compute gets faster. Resolving dimensions and addresses on one queue would serialize work that could run concurrently.

Patching is not the only host work involved. Preparing a tensor in the device’s expected layout can also run alongside device compute and be connected to its consumer by an event.

These events are software signals checked by the scheduler. They are single-use, so a fresh event backs each edge. The interface can later use hardware event support without changing the layers above it.

What the mapping buys, and what it costs

User experience

The surface a user touches is the one they already know. tensor.to("spyre"), torch.compile, and torch.Stream use the same PyTorch concepts as other devices. That does not automatically provide operator coverage, profiling, distributed collectives, CUDA-specific extensions, or model-specific kernel quality; those remain backend work.

For compile-backed eager operations, eager mode is not a second implementation. Torch-spyre registers a PrivateUse1 kernel around the same decomposition used during compilation. Inside a torch.compile trace, the wrapper calls that decomposition directly so PyTorch can capture it into the surrounding FX graph. Outside a trace, the first eager call compiles and caches the decomposition; later calls reuse the compiled entry point. These operations therefore reuse the same Inductor and launch machinery instead of maintaining a handwritten runtime beside it.

That choice still exposes first-use compilation; unsupported operations can fail or fall back, and different tensor extents can require another compiled artifact. The benefit is one execution path whose fixes and regressions apply to both modes.

A compile-backed eager operation is also the smallest possible compiled program: one operation, real tensors, real device memory, and a real launch. If a model fails but the corresponding eager operation passes, the fault is more likely in composition than in the operation. A test for a single operation exercises the whole compile-and-launch stack, so a per-operation suite doubles as an integration suite.

And because residency is generic rather than hand-written per case, tensors stay on the device without bespoke residency code for each kind of tensor. That is what makes eager mode practical.

When the device behaves like a normal PyTorch device, a new model implementation is mostly just a model implementation. It can be picked up and run rather than ported, because its operations already dispatch and its tensors already live in the right place. The same holds one level up. Inference techniques that restructure how a model is driven – speculative decoding, say, where a small draft model proposes tokens a larger one then verifies – are written against streams, tensor residency and ordinary control flow. If those are the framework’s own abstractions, the technique can be tried on the device without anything being added to the compiler.

The compiler makes the arithmetic fast, but much of an inference stack’s performance comes from how work is scheduled, batched, overlapped and speculated. A backend that maps cleanly onto the framework’s primitives leaves that space open to people who never touch the compiler.

Performance

The performance argument is about what leaves the launch path. Submitting typed operations to a queue, rather than constructing a graph to describe them, removes a disk read, a deserialize, a per-argument stitch, and a compile-and-parse from every single launch. A launch becomes building a few operations and enqueueing them.

On a real model the effect compounds, because those costs were being paid per launch. On Granite 3.3 8B at batch size 1 and sequence length 1024, taking graph construction off the transfer path alone is worth 9.9% on prefill and 13.6% on decode – for a change that touches only how bytes move. Taking it off compute dispatch as well brings the totals to 1.7x faster prefill and 2.4x faster decode, with roughly three quarters of the gain coming from the compute path.

The concurrency described in Mapping 4 is a separate opportunity: every piece of host preparation moved off the compute path reduces device idle time.

What PyTorch does not provide by itself

PyTorch has a device-event abstraction, but the cross-stream ordering described above has to be enforced where the operations actually sit – in the runtime scheduler, which is the layer that sees their device-memory footprints. The runtime therefore builds software event primitives and derived dependencies first. Exposing equivalent user-created events through PyTorch’s public event hooks is separate integration work.

What we put back

Building on a framework’s extension points is not a one-way transaction. PyTorch’s core developers maintain openreg, the upstream reference backend and accelerator simulator that shows how an out-of-tree device can register through PrivateUse1 without modifying core PyTorch, with contributions from hardware ecosystem partners. Contributing improvements there is how backend work like torch-spyre shortens the path for the next device team, and participating in the PyTorch Technical Advisory Council is how the constraints of devices like this one reach the people designing the next extension point. Our integration with PyTorch’s Cross-Repository CI Relay (CRCR) tests upstream changes against Spyre and reports results to PyTorch’s CI dashboard, helping catch compatibility regressions as the framework evolves. We also built reusable tools for selecting and adapting upstream tests, which we plan to contribute upstream so other out-of-tree accelerators can benefit. A backend that only consumes an interface leaves the interface no better than it found it.

Conclusion

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.

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.

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.

The post Building Spyre as a Native PyTorch Device appeared first on PyTorch.