Fuse ONNX Attention

One input, one output, no settings. This node takes the ONNX graph you just exported and rewrites its self-attention blocks into ONNX Runtime's fused attention operator, writing <name>_fused.onnx next to the source and handing back a handle you can run immediately. The README files it as "optional, ~faster," which undersells the decision - it's the difference between an ONNX model that's merely portable and one worth running.

Why it exists

A naively traced transformer exports attention as the literal chain of ops from the PyTorch code: separate projections, a matmul, a scale, a softmax, another matmul, a reshape or two. Numerically fine, and ONNX Runtime's graph optimizer does what it can, but it will not invent a fused kernel out of that soup. ORT ships one - a flash-attention-style operator in the com.microsoft domain - and the only way to get it into your graph is to rewrite the graph. That's the whole job: each of the encoder's attention subgraphs gets replaced by a single fused node.

The node path uses the attention variant of the fuser, which also folds the QKV projection into the fused op, rather than the more conservative mha. More fusion, fewer nodes, more speedup.

How it works

It does not blindly pattern-match. Before rewriting anything, it validates that the graph really is the DA2 layout at the encoder it claims to be - the block count, head count, and hidden size all have to line up (Small is 12 blocks / 6 heads / 384 wide, Large is 24 / 16 / 1024, Giant 40 / 24 / 1536). If the graph doesn't match, it fails instead of emitting something subtly wrong. The encoder is read from the graph's own metadata, which the exporter stamped in, so you're not asked to specify anything.

The fused model keeps a depthaccel.attention metadata entry recording which operator it was rewritten to, which is how the loader can tell what it's holding.

Inputs and output

onnx_model in, onnx_model out - both DEPTHACCEL_ONNX_MODEL. That's it, no precision or device knobs, because fusion is a pure graph transformation.

The output replaces the input at the point where you'd otherwise run Estimate Depth (ONNX). And because this node is also a graph endpoint (is_output_node), you can queue [Export to ONNX] → [Fuse ONNX Attention] with nothing attached and it will just write the _fused.onnx file for later. After a restart, Load Depth Model (ONNX) will see it in the dropdown - and will actually prefer it, since the loader's default pick is the fused file when one exists.

The one thing to know: fused is ORT-only

The fused graph uses com.microsoft operators. ONNX Runtime understands them and TensorRT does not. So if your plan is TensorRT, keep the explicit graph and build the engine from that - the engine builder refuses a fused input with a clear error rather than producing a bad engine, which is the right behavior and still a detour you'd rather not take. Practical shape of it: export once, fuse if you want a fast ONNX handle, build the engine from the unfused file. Two files, two purposes, both cheap on disk.

Install

Fusing needs the onnx package (the operation is done with onnx plus NumPy - no ONNX Runtime involvement in the rewrite itself). Returning the handle needs onnxruntime-gpu, because the node loads the file it just wrote before handing it over:

python -m pip install onnx onnxslim onnxruntime-gpu

Both are commented out in the pack's requirements.txt, so ComfyUI Manager won't install them for you. Uncomment and reinstall, or just pip them into the same venv - and if you skip ONNX Runtime, expect the write to succeed and the node to fail right after with a message naming the package, which at least leaves you with the artifact.

Where it fits

[Load Depth Model (PyTorch)] → [Export to ONNX] → [Fuse ONNX Attention] → [Estimate Depth (ONNX)]

That's the fast no-engine pipeline, and the pack's own numbers put fused ONNX at roughly 1.4× PyTorch and about 1.4× slower than TensorRT. Skip this node only when you're going to TensorRT instead - in which case the fused file is dead weight, and you should delete it before it becomes the default pick in the ONNX loader and confuses you six weeks later.