Skip to content

ONNX export: fused GELU operator in the matcher - #6

Merged
sarlinpe merged 1 commit into
sarlinpe/onnx-exportfrom
sarlinpe/onnx-fused-gelu
Oct 30, 2026
Merged

sarlinpe merged 1 commit into
sarlinpe/onnx-exportfrom
sarlinpe/onnx-fused-gelu

Conversation

@sarlinpe

@sarlinpe sarlinpe commented Oct 28, 2026 •

Copy link
Copy Markdown
Member

Stacked on #5. This PR exports the GELU activations of the matcher as ONNX Runtime's com.microsoft::Gelu operator (exact, erf-based) instead of Div / Erf / Add / Mul / Mul.

Why

  • ONNX Runtime doesn't fuse the exported GELU pattern (Div(√2) → Erf → Add(1) → Mul(0.5) → Mul(x)). Each of the 36 GELUs per pair therefore runs as five element-wise kernels over a 4096×512 tensor.
  • The exporter swaps the nn.GELU modules of the exported copy of the matcher for a custom op, which the translation table maps to the fused operator (as for the attention). --no_fused_attention disables both, for standard ONNX operators.
  • Only the matcher changes; the descriptor model is unchanged.

Results

GPU time per pair: ONNX Runtime 1.23, A100, N = 2048, CUPTI kernel profile, device-resident I/O.

fp16 matcher GPU time Kernels
Current export 30.42 ms 1470
With fused GELU 29.58 ms (−0.84) 1326

Disabling ONNX Runtime's MatMulAddFusion would save another ~1.4 ms on this setup. That was not adopted, because it depends on cuBLAS's kernel choice; see colmap/colmap#4914.

Parity: python onnx/check.py passes (N = 512 / 2048 / 4096, min_covisibility 0 / 0.8):

  • fp16: at most 2 differing matches.
  • fp32: at most 1 differing match on CUDA, previously exact. ONNX Runtime's Gelu kernel rounds slightly differently than F.gelu.
  • The covisibility score is within 1e-3 of PyTorch.

@Phil26AT Phil26AT left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, nice find!

Export the GELU activations of the matcher as ONNX Runtime's com.microsoft::Gelu
operator instead of Div/Erf/Add/Mul/Mul, which ONNX Runtime does not fuse, by
swapping the nn.GELU modules of the exported copy of the matcher. This saves
about 0.9 ms of GPU time per pair with the float16 matcher.
--no_fused_attention also disables it.
@sarlinpe
sarlinpe force-pushed the sarlinpe/onnx-fused-gelu branch from 6f99534 to 212c0c8 Compare October 30, 2026 02:26
@sarlinpe
sarlinpe merged commit a8e9457 into main Oct 30, 2026
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants