Making PyTorch skip the part that didn’t change

HS‑002 2026‑08‑07 PyTorch extension

A lot of simulation and streaming code has the same shape. There is a big sparse matrix that describes how things are wired together — a grid, a mesh, a circuit, a graph — and it does not change. There is a state vector that does change, but only a little: a few hundred entries out of a hundred thousand, every step. And then there is a line that multiplies the two together, over and over, for as long as the simulation runs.

In PyTorch that line is torch.sparse.mm(A, x), and it does the whole multiply every time. It has no way to know that almost nothing moved. We built a PyTorch extension that does know, and this note reports what it is worth — including the cases where it is worth nothing and the ones where it costs you.

The idea in one paragraph

If A is fixed, then multiplying it by a slightly different vector gives a slightly different answer, and you can compute the difference directly: A(x + dx) = Ax + A·dx. You already have Ax from last step. So the only work left is A·dx, and if dx is zero almost everywhere, that touches only the columns of A where it isn’t. Everyone in this area knows that much. What we add is the layout: we cut A into tiles ahead of time and record which tiles each column falls in, so at run time we can walk only the tiles the changed columns pass through and never look at the rest. The answer is the same one PyTorch would have given — agreement to 1.6e‑15 relative across the whole sweep below, which is floating-point noise, not an approximation we chose.

What you actually install

An out-of-tree extension. You do not build PyTorch from source and you do not patch it: the ops register through TORCH_LIBRARY like any other custom op, so they are traceable and they compose with the rest of your model. There is a thin Python wrapper over the top for the common case. The apply writes into a tensor you own, in place, which is a small thing that turns out to matter a great deal — see the ladder below.

op   = HSTOperator(A, tile=8)          # compile the tile index once
sess = op.session(cols, B=8)           # schedule for the columns you expect

for step in range(n):
    cols, dx = advance(...)            # your solver, unchanged
    sess.apply_cols(cols, dx, y)       # y += A[:, cols] @ dx, in place

The schedule is compiled once for a set of column tiles and then reused: apply_cols takes a different column set each step, which is legal as long as the columns stay inside tiles the schedule already covers. That is what lets one schedule survive a changed set that keeps moving, and it is the mechanism behind everything below. That is also the whole bargain: you pay a preparation cost up front so that every later step can be cheap, which means the method is worth having when the loop is long and worth nothing when it is short.

How much faster — and than what

This is the part that needs care, because you can get almost any number you like out of this comparison depending on what you put on the other side. So here is the whole ladder, each rung measured in the same run on the same five operators, at a batch width of 8.

Read that from the top. Recomputing in PyTorch is the starting point. Slicing the changed columns out once and multiplying only the slice — which is the best you can do without leaving PyTorch — is 3.2× faster than that. Writing the same delta by hand in C++ is another 6.8× on top. And our scheduled version is another 2.5× on top of that, for 54× end to end.

The honest reading of that chart is that the big number is mostly not ours. Anyone who swaps torch.sparse.mm for our extension will measure 16.7×, and that is a real thing that will happen to them — but it factors as 6.8× for leaving PyTorch and 2.5× for the schedule, and only the second number is ours. torch.sparse.mm has no out=, so every call allocates a fresh output block, and its cost tracks the number of rows rather than the size of the slice you asked for. Those are structural facts about the API, not benchmark artefacts, and any competent C++ delta collects most of that speedup without us.

So the number we care about internally is the last rung only: what is the schedule worth against a hand-written delta that is already doing the right thing? That comparison is the rest of this note.

Against a delta that is already good

The baseline here is deliberately strong — a compressed row list compiled for the exact set of changed columns, threaded, with the integer widths and restrict qualifiers you would use if you cared. It is what a good engineer produces after a week on the problem. Against that baseline the answer depends almost entirely on how the changed set moves, and much less on how fast it changes.

local drift — the changed set walks along the operator uniform jump — it teleports somewhere unrelated

When the changed set drifts — a wavefront moving across a mesh, a load moving along a network, the ordinary behaviour of a time-stepping simulation — we run 6.9× the baseline at the lowest churn rate and still 4.8× at the highest rate we have measured. All 21 operators are faster, at every churn rate in that range. The reason is the tiling: when the set moves a little, most of the tiles it needs are the ones it needed last step, so we pay only for the new ones. The baseline is compiled for an exact column set, so any movement at all forces it to recompile.

When the changed set jumps — a genuinely random new set every step, nothing reused — the advantage collapses from 2.9× to 1.4× and the number of operators we win on falls from 20 of 21 to 12 of 21. That is the mechanism working exactly as described: there is no locality left to exploit, so there is nothing for the schedule to save.

And when nothing moves at all — a fixed set of changed columns, step after step — we lose, at 0.48×. This is not a defect and it is not going to be fixed. Tiles are a fixed grid, so a tile containing one changed column still gets scanned as a tile; the baseline touches exactly the entries it needs. Our best case on a frozen set is a tie, and padding cannot be negative. If your changed set never moves, use the baseline.

Neither arm wins everywhere, so we ship both

Put the two motion models side by side per operator and the picture stops being a single number. Under drift every operator is on our side, by between 1.3× and 34×. Under jump at the fastest churn rate measured, nine of the same 21 operators are on the baseline’s.

local drift, ρ=0.002 uniform jump, ρ=0.25

You cannot pick between the two by looking at the matrix, and the clearest evidence is that all nine operators the baseline wins under jump are ones we win under drift at the same churn rate, by between 1.15× and 6.6×. Same matrices, same rate, opposite answers. We also tried the density statistic that ought to separate the two classes: it overlaps badly, with our wins running as low as 0.18 and the baseline’s as high as 0.80. What decides is how the workload moves, and that is not a property of the matrix you can read off in advance.

So the product does not choose by a rule. It runs both arms briefly on the caller’s own data, keeps the faster one, and re-probes when the workload shifts underneath it. That is worth 5.1× in steady state under drift, or 4.0× once the cost of the probing itself is charged against it, and 1.9× / 1.5× under jump. The reason it is worth doing at all is the size of a wrong answer: the worst single case we have measured costs 35.6× if you guess wrong in one direction, and 2.4× in the other.

Whether it will help you

Six things have to be true, and they fail independently. The cheapest ones to check are first.

• The sparse operator is fixed. Same nonzeros, same values, across many steps. If you rebuild it every iteration because it depends on the current estimate, there is nothing to amortise.
• The step is linear. A softmax, a ratio, a percentile or a top-k anywhere in the loop and the difference of the outputs is no longer the output of the difference.
• The change is sparse in the operator’s own dimension. A few dirty columns, not a few dirty rows of a batch.
• The dirty set is localised. It may move — drift is the best case we have — but it should not teleport.
• It runs in your process. The saving is microseconds per step, so a network hop or a serialisation boundary erases it before you see it.
• The hot loop is a multiply, not a solve. A fixed operator whose bottleneck is a factorisation or a triangular substitution does not decompose this way, even with the other five perfect. This one catches people, because everything else looks right.

A useful shortcut, if you want to know in one question: is the loop you want to speed up the time-stepping loop, or the Newton loop? Traced against real solves, transient propagation and moving loads run 145 to 190 steps with a dirty set that barely turns over, which is the drift case above. Newton iterations re-linearise, so the operator is not fixed across them, and they converge in a median of ten steps — too few to amortise anything. The boundary runs straight through the middle of a single simulation.

← All notes