FuseFlow: A Fusion-Centric Compilation Framework for Sparse Deep Learning on Streaming Dataflow

Rubens Lacouture, Nathan Zhang, Ritvik Sharma, Marco Siracusa, Fredrik Kjolstad, Kunle Olukotun and 1 other

Sparse deep-learning models skip the zeros in weights or activations to save work, and streaming dataflow chips are great at running them by wiring compute units directly together — but to hit that hardware you need a compiler that fuses the operations so intermediate results never touch memory, and prior sparse compilers could only fuse within a single expression; this paper (Stanford — Rubens Lacouture, Fredrik Kjolstad, Kunle Olukotun, Olivia Hsu, and colleagues) builds FuseFlow, the first compiler to do general cross-expression and partial fusion of sparse operations from PyTorch to dataflow graphs, and shows that full fusion is not always best — the right granularity depends on the model — reaching about 2.7 times speedup on GPT-3 with BigBird block-sparse attention.

The argument in one glance

Problem

Sparse models fit dataflow chips beautifully, if you can program them

Many deep-learning models are sparse: their weights, activations, or inputs are full of zeros, whether naturally or because they were pruned. Skipping those zeros cuts the compute and memory a model needs, but it also makes memory access irregular and hard for conventional hardware. A promising answer is streaming-dataflow hardware — reconfigurable dataflow architectures that connect coarse-grained compute units directly to one another, so data streams from one unit to the next without going through expensive caches and local memories. This style is a natural fit for sparse computation. The problem is getting your model onto such a chip. To run efficiently, the operations of the model have to be fused together so that the intermediate results flowing between them are streamed unit-to-unit and never written out to memory, and expressing that fusion is exactly what has been hard.
What people did

Use a sparse tensor-algebra compiler for dataflow

The prior state of the art, a compiler called Custard built on the Sparse Abstract Machine dataflow model, was the first to compile sparse tensor-algebra expressions to streaming dataflow graphs. You write a sparse expression, it generates a fused dataflow graph for the hardware, and it lets you choose the order in which the computation iterates.
Why it fell short

It fuses only within one expression, and only all-or-nothing

That compiler has two limits that make it unsuitable for real machine-learning models. First, it compiles one expression at a time, so it cannot fuse operations across expressions — but a model is many expressions, and cross-expression fusion is the whole point of a machine-learning compiler. Second, the fusion it does within an expression is total: it fully fuses everything, with no way to fuse partially. Yet partial fusion is often what you want, because the fusion-versus-recomputation tradeoff is fundamental — fusing eliminates intermediate tensors and cuts memory traffic, but fusing too much forces values to be recomputed, and fusing too little dumps intermediates to memory. For sparse models this tradeoff is even sharper, since fusion can change the asymptotic complexity of the computation.
This paper's idea

Compile whole sparse models with cross-expression and partial fusion, as a schedule

FuseFlow is a compiler that takes a sparse machine-learning model written in PyTorch and lowers it to a fused sparse dataflow graph for reconfigurable dataflow hardware, and it is the first to support general cross-expression fusion of sparse operations — and partial fusion too. It compiles complete inference pipelines, including the non-linear operations and masking that real models need, not just isolated tensor-algebra kernels. It exposes fusion granularity as a user schedule, so a developer can dial the amount of fusion up or down and explore the tradeoff per model, and it uses two new intermediate representations — a fused Einsum form that tracks coordinate and value streams across fused expressions, and a fusion table that names and memoizes intermediate streams so the compiler can emit the fused graph efficiently. Exploring the design space across four real sparse models, FuseFlow shows full fusion is not always optimal — the best granularity depends on the model — and a heuristic prunes the bad configurations. It reaches about 2.7 times speedup over an unfused baseline on GPT-3 with BigBird block-sparse attention.

Why this problem is worth solving

The hardware trend and the model trend point at the same place, and only the compiler is missing.

Sparse deep learning means running models whose weights, activations, or inputs contain many zeros — from pruning, from natural structure, or from induced sparsity patterns — and skipping the zero entries during computation. This saves compute and memory, which matters more as models scale, but it produces irregular, data-dependent memory access that general-purpose hardware handles poorly. In response, architects build specialized accelerators, and increasingly they use dataflow: instead of a processor repeatedly loading operands from caches and writing results back, a dataflow architecture wires coarse-grained functional units directly together so data streams from one to the next. These reconfigurable dataflow architectures lean less on expensive memory operations, which suits sparse computation especially well, because the value that matters is kept moving rather than staged in and out of memory. So both sides are aligned — sparse models want to avoid work, dataflow hardware wants to avoid memory traffic — and the only thing standing between a model expressed in a framework like PyTorch and this hardware is a compiler that can bridge them. That is the gap FuseFlow fills, and it is worth filling because without it, the efficiency these architectures promise for sparse models is simply not reachable from the tools practitioners actually use.

The heart of the difficulty is fusion, and for sparse models fusion is a richer, more consequential decision than it is for dense ones.

Fusion means merging operations so the intermediate tensor between them is never written to memory — the output of one operation streams directly into the next as it is produced. On dataflow hardware this is essential, because the intermediates that fusion eliminates would otherwise have to be materialized and moved, which is exactly the memory traffic the architecture exists to avoid. But fusion is not free, and the tradeoff is fundamental: fuse two operations and you save the intermediate, but if a later operation needs a value more than once, fusion can force you to recompute it rather than read it back from a stored tensor; fuse too little and you pay to write and re-read all the intermediates. The right point on this spectrum is partial fusion — fuse enough to cut memory traffic but not so much that recomputation dominates — which is precisely the insight behind FlashAttention. For sparse computation the stakes are higher, because how you fuse can change the asymptotic complexity of the computation, not just its constant factors, so the wrong choice is not merely slower but algorithmically worse. This is why a sparse machine-learning compiler must not bake in one fusion strategy but expose fusion granularity as a knob to be tuned per model, and why a tool that can only fully fuse a single expression is inadequate.

What people did before, and where it broke

Dense tensor compilers with cross-kernel fusion. Machine-learning compilers routinely fuse operations across kernels for dense models to cut memory traffic. Where it broke. They target dense computation and do not handle the irregular, data-dependent structure of sparse tensors, so their fusion machinery does not transfer to sparse models on dataflow hardware.
The Sparse Abstract Machine and the Custard compiler. Custard was the first compiler to lower sparse tensor-algebra expressions to streaming dataflow graphs on the Sparse Abstract Machine, generating fused graphs and letting users pick the iteration order. Where it broke. It compiles individual expressions, so it cannot fuse across expressions — the key capability a model needs — and its intra-expression fusion is total, with no partial fusion.
Full fusion of everything. Fuse the entire computation end to end to eliminate all intermediate tensors. Where it broke. For sparse models full fusion is not always optimal; fusing too aggressively forces recomputation, and the paper finds the best granularity depends on the specific model.
Hand-written rewrites to force some cross-expression reuse. Manually restructure the computation to get a bit of fusion across expressions with existing tools. Where it broke. It takes substantial user effort and is fragile; the paper shows a hand-written Custard-plus-Stardust rewrite reaches only part of the speedup FuseFlow gets automatically with less effort.
The gap. No compiler could take a whole sparse model and perform general cross-expression fusion and partial fusion of its sparse operations, exposing fusion granularity as a schedule so the fusion-recomputation tradeoff could be explored per model and lowered efficiently to a streaming-dataflow graph.

How FuseFlow works

From a sparse PyTorch model to a tunably-fused streaming dataflow graph

↓ →
new
↓ →
new
↓ →
new
↓ →
new
↓ →
new

The idea

Make fusion a tunable knob, not a fixed strategy, because the sweet spot moves with the model. The conceptual center of FuseFlow is the recognition that there is no single right amount of fusion for sparse models. Fusing merges operations to eliminate the intermediate tensors that would otherwise clog memory, but fuse too much and you force recomputation of values a later step reuses; fuse too little and you pay to materialize and re-read intermediates. Prior sparse compilers baked in a choice — full fusion of a single expression — which cannot express the partial fusion that is often optimal, the very thing FlashAttention showed for attention. FuseFlow instead exposes fusion granularity as a user schedule, letting the developer mark exactly which operations to fuse and how the computation iterates, so the fusion-recomputation tradeoff can be explored rather than assumed. This matters more for sparse than dense computation because fusion can change the asymptotic complexity, so the choice is algorithmic, not just a constant-factor tuning. Turning a hard-wired optimization into an exposed schedule the programmer can dial, so a fusion-and-reuse tradeoff can be searched per workload, is the same design philosophy the composable sparse-compilation-abstractions work covered elsewhere on this site is built on, decomposing sparse compilation into choices a user can compose rather than a fixed pipeline. FuseFlow adds the finding that this exploration is necessary, because full fusion is not always best.

Track coordinates and values across expressions, then name the streams so you can weave them into one graph. The reason cross-expression fusion of sparse operations was missing is that it is genuinely hard to represent: a sparse computation carries not just values but the coordinates that say where the nonzeros are, and to fuse operations across expression boundaries you must track how both flow through the merged computation. FuseFlow introduces a fused Einsum representation for exactly this — an Einstein-summation form that follows the indexing data and the values across the fused expressions. But representing the fusion is only half the job; lowering several fused expressions into a single streaming dataflow graph is where the complexity explodes, because subgraphs are shared and referenced before they are materialized. FuseFlow’s second representation, the fusion table, solves this by naming and memoizing the intermediate streams, organized by iteration, tensor, and node, so the compiler can point at a subgraph before it exists and emit the woven-together graph efficiently. Designing dedicated intermediate representations that carry Einstein-summation semantics down to a dataflow substrate, so high-level sparse expressions become concrete streaming hardware graphs, is the same compilation strategy the work on compiling Einstein summations to locality-aware dataflow covered elsewhere on this site pursues, using the Einsum abstraction as the bridge from math to a dataflow mapping.

Build on a streaming-dataflow abstract machine, and make it speak machine learning. FuseFlow does not invent its target from scratch; it stands on the Sparse Abstract Machine, a dataflow abstract machine whose handful of composable blocks can express any sparse tensor-algebra expression, support fusion across expressions, and allow different iteration orderings, and which compiles down to real dataflow hardware. That makes it a natural intermediate representation for targeting sparse hardware — but it was built for tensor algebra, not for whole machine-learning models. So FuseFlow extends it into an ML-capable form: it adds dense blocks so block-sparse tensors (which mix dense sub-blocks with sparse structure, as in BigBird attention) can be handled, plus non-linear functions and masking operations that no pure tensor-algebra machine had, and it lets the generated graphs run data-parallel in addition to the pipeline parallelism dataflow gives for free. Taking a principled streaming-dataflow abstraction for sparse tensor algebra and elaborating it into the substrate a full sparse model can compile to is the same foundational role the Sparse Abstract Machine work covered elsewhere on this site plays, providing the composable dataflow vocabulary on which a compiler like FuseFlow can then build model-level fusion.

Following a fused attention model through FuseFlow

Compiling GPT-3 with BigBird block-sparse attention to a fused dataflow graph

A trace, drawn from the paper's design and evaluation, of how FuseFlow lowers a sparse model and why fusion granularity matters.

Step 1 of 5

A block-sparse attention model arrives from PyTorch

The input is GPT-3 Small with BigBird attention, whose attention matrix is block-sparse — most blocks are skipped, a few dense blocks are kept. It comes in as PyTorch with sparse annotations, including the non-linear softmax and the masking that attention requires, which pure tensor-algebra compilers could not express.

FuseFlow's added support for dense blocks, non-linear functions, and masking is what lets this whole model be compiled, not just its matrix multiplies.

Why it works

Cross-expression fusion removes intermediates a per-expression compiler must materialize. Compiling one expression at a time forces every inter-expression tensor to memory → but fusing across expressions streams those results unit-to-unit instead → so the memory traffic that dominates on dataflow hardware is eliminated where it helps.
Partial fusion finds the point full fusion overshoots. Fully fusing everything can force recomputation that outweighs the memory saved → but exposing fusion granularity as a schedule lets the compiler stop short of that → so each model runs at its own optimum, which the paper shows is often not full fusion.
Tracking coordinates and values makes sparse fusion representable. Sparse computation carries nonzero locations, not just values, so naive fusion loses the structure → but the fused Einsum form follows both coordinate and value streams across expressions → so operations on sparse tensors can be fused end to end correctly.
Naming and memoizing streams tames the lowering. Weaving many fused expressions into one graph means referencing subgraphs before they exist → but the fusion table names and memoizes the intermediate streams → so the compiler emits the fused dataflow graph efficiently instead of drowning in materialization order.
A heuristic keeps design-space exploration cheap. The number of equivalent fusion configurations is large and only some are good → but a fusion heuristic prunes the inefficient ones using lightweight analysis → so finding the granularity that yields the speedup does not require exhaustively simulating every schedule.

What it costs (trade-offs and scope)

Fusion granularity is a schedule the user must choose or search. FuseFlow exposes the tradeoff rather than deciding it automatically, so getting the best performance means exploring configurations; the heuristic prunes bad ones, but the developer is still in the loop marking fusion regions and iteration orders, which is more involved than a push-button compiler.
The sparse data-structure type must be fixed before compilation. FuseFlow supports sparse tensors from any source, but each tensor’s format has to be determined at compile time, so it does not adapt to sparsity structure discovered only at runtime.
It targets a specific hardware model and is evaluated largely in simulation. The target is streaming-dataflow / reconfigurable dataflow architectures via the Sparse Abstract Machine, and the microarchitectural results come from a cycle-accurate simulator (with FPGA-validated agreement); the numbers reflect that setting, not arbitrary hardware, and full deployment breadth on many fabricated accelerators is not the focus.
The gains and the “not always full fusion” finding are model-dependent. The design-space exploration shows speedups from roughly 1.5 to 3.9 times across configurations and the about-2.7-times headline for GPT-3 with BigBird; these are for four specific sparse models (a sparse autoencoder, two graph networks, and a block-sparse transformer) at particular sparsity levels, so other models may land differently.
It is a research compiler with real dependencies and effort. FuseFlow is the first academic end-to-end sparse-ML compiler for these architectures, implemented in MLIR; adding a new model, format, or operation still requires engineering within that framework, and the value is in the fusion methodology, not a turnkey production toolchain.

The evaluation

Compiling sparse deep-learning models to streaming dataflow

Property	Custard (per-expression)	Full fusion always	FuseFlow
Fuses across expressions
Supports partial fusion
Handles whole ML models
Finds model-specific optimum
Reported outcome

A comparison drawn from the paper. FuseFlow is evaluated on four sparse ML model classes (a sparse autoencoder, GCN, GraphSAGE, and GPT-3 Small with BigBird block-sparse attention) targeting a cycle-accurate dataflow simulator with FPGA-validated agreement; the about-2.7-times figure is the paper's reported speedup over an unfused baseline for GPT-3 with BigBird.Underlined values explain themselves when clicked.

Where this stops working

What to keep in mind

The full paper text was available (the authors’ arXiv version, 2511.04768), so this account of the mechanism — sparse deep learning on reconfigurable dataflow architectures, the Sparse Abstract Machine target and the prior Custard compiler’s per-expression and full-fusion limits, the fusion-recomputation tradeoff and why it is sharper for sparse computation, FuseFlow’s cross-expression and partial fusion exposed as a schedule, the fused Einsum and fusion-table intermediate representations, the ML extensions (dense blocks, non-linear functions, masking) and data parallelism, and the design-space exploration with a pruning heuristic — follows the paper’s own description. The headline numbers are the paper’s reported results: about 2.7 times speedup over an unfused baseline for GPT-3 with BigBird block-sparse attention, with speedups ranging roughly 1.5 to 3.9 times across the explored configurations, evaluated on four sparse model classes on a cycle-accurate dataflow simulator validated against FPGA. Those figures are for specific models and sparsity levels on a particular hardware model; results vary with the model and schedule. Exact per-configuration numbers, the IR algorithms, and the heuristic’s details should be read from the paper itself.

Background

Why fusion is the recurring lever in ML compilation

It helps to see that fusion is one of the most important ideas in the whole field of compiling machine-learning models, dense or sparse. A modern model is a long chain of operations, and the naive way to run it — execute each operation, write its output to memory, read it back for the next — spends enormous time and energy shuttling intermediate tensors in and out of memory rather than computing. Fusion attacks this directly by keeping intermediates in fast on-chip storage or, on dataflow hardware, streaming them straight from one compute unit to the next so they never hit memory at all. The reason it keeps reappearing across compilers and accelerators is that memory movement, not arithmetic, is usually the bottleneck for these memory-bound workloads, and fusion is the most direct way to cut it. The subtlety FuseFlow adds is that more fusion is not always better, because past a point it forces recomputation, so the real art is choosing the granularity — a search, not a default. That centrality of restructuring a computation so intermediate results avoid a round trip to memory is the same lever the optimized fused-attention-dataflow work covered elsewhere on this site pulls, fusing an attention pipeline so its large intermediate scores are never materialized. FuseFlow generalizes that hand-crafted instinct into a compiler that can search for the right fusion across any sparse model.

If you remember only three things

1. Sparse deep-learning models skip zeros to cut compute and memory, and streaming-dataflow hardware (reconfigurable dataflow architectures that wire compute units directly together) is well-suited to run them, but reaching that hardware requires fusing operations so intermediate tensors are never materialized to memory — and prior sparse tensor-algebra compilers for dataflow (Custard on the Sparse Abstract Machine) could only fuse within a single expression and only fully, unable to fuse across expressions or partially, even though the fusion-recomputation tradeoff is fundamental and, for sparse models, can change asymptotic complexity.
2. FuseFlow is the first compiler to support general cross-expression and partial fusion of sparse operations, lowering complete sparse PyTorch models (including non-linear ops and masking) to fused dataflow graphs and exposing fusion granularity as a user schedule so the tradeoff can be explored. It introduces two IRs — a fused Einsum form that tracks coordinate and value streams across fused expressions, and a fusion table that names and memoizes intermediate streams to emit the graph efficiently — extends the Sparse Abstract Machine with dense blocks, non-linear functions, and masking, and adds data parallelism alongside dataflow's native pipeline parallelism.
3. Across four real sparse models (a sparse autoencoder, GCN, GraphSAGE, and GPT-3 Small with BigBird block-sparse attention), FuseFlow's design-space exploration shows full fusion is not always optimal — the best granularity depends on the model — with speedups ranging roughly 1.5 to 3.9 times across configurations and about 2.7 times over an unfused baseline on GPT-3 with BigBird; a heuristic prunes inefficient configurations cheaply. It targets a cycle-accurate dataflow simulator (FPGA-validated), requires the sparse format to be fixed at compile time, and leaves fusion-region selection as a schedule the developer explores.

Check yourself

Check your understanding

1. Why is fusion essential for running sparse deep-learning models on streaming-dataflow hardware?
2. What were the two key limitations of the prior sparse tensor-algebra compiler (Custard) that FuseFlow removes?
3. How does FuseFlow make cross-expression fusion of sparse operations both expressible and efficient to compile?
4. What did FuseFlow's design-space exploration reveal, and what were its results?

Sources

The paper is “FuseFlow: A Fusion-Centric Compilation Framework for Sparse Deep Learning on Streaming Dataflow” by Rubens Lacouture, Ritvik Sharma, Fredrik Kjolstad, Kunle Olukotun, and Olivia Hsu (Stanford University; Hsu also Carnegie Mellon University), Nathan Zhang (SambaNova Systems), and Marco Siracusa (Barcelona Supercomputing Center), published at ASPLOS 2026, DOI 10.1145/3779212.3790165. It is licensed CC BY.

This explainer is based on the paper’s full text, obtained from the authors’ arXiv preprint (arXiv:2511.04768). The account of the mechanism — sparse deep learning and its fit with reconfigurable dataflow architectures; the Sparse Abstract Machine as a streaming-dataflow IR and the prior Custard compiler’s limitation to per-expression, full fusion; the fundamental fusion-recomputation tradeoff and why sparsity sharpens it (up to changing asymptotic complexity); FuseFlow’s general cross-expression fusion and partial fusion exposed through a scheduling language; the two new intermediate representations (a fused Einsum form tracking coordinate and value streams across expressions, and a fusion table that names and memoizes intermediate streams for efficient lowering); the ML extensions to the Sparse Abstract Machine (dense blocks for block-sparse tensors, non-linear functions, and masking) plus data parallelism alongside native pipeline parallelism; and the design-space exploration with a pruning heuristic — follows the paper’s own description. The reported results cited here are the paper’s own: about a 2.7 times speedup over an unfused baseline for GPT-3 Small with BigBird block-sparse attention, speedups ranging roughly 1.5 to 3.9 times across explored configurations, and the finding that full fusion is not always optimal, evaluated on four sparse model classes (a sparse autoencoder, a graph convolutional network, GraphSAGE, and GPT-3 with BigBird) on a cycle-accurate dataflow simulator with validated FPGA agreement, with FuseFlow implemented in MLIR. Exact per-configuration numbers, the IR algorithms, and the heuristic’s details should be read from the paper itself.