Skip to content

FlexGlue: disable MatMulAddFusion on CUDA (faster Linear layers) - #4914

Closed
sarlinpe wants to merge 1 commit into
sarlinpe/flexgluefrom
sarlinpe/flexglue-bias-epilogue
Closed

sarlinpe wants to merge 1 commit into
sarlinpe/flexgluefrom
sarlinpe/flexglue-bias-epilogue

Conversation

@sarlinpe

Copy link
Copy Markdown
Member

Stacked on #4908 (FlexGlue matcher) and #4907.

Summary

ONNX Runtime's MatMulAddFusion turns each Linear layer of the FlexGlue matcher (MatMul + bias Add) into a Gemm. Its CUDA kernel first broadcasts the bias into the output with an extra rank-1 GEMM (k = 1), which cuBLAS runs on a slow non-tensor-core kernel (ampere_fp16_sgemm_fp16_128x128_nn, 162 launches per pair, 2.2 ms of GPU time), and only then runs the real tensor-core GEMM with beta = 1.

This PR disables the fusion for the FlexGlue matcher on CUDA. The Linear layers then run as a tensor-core MatMul (beta = 0) plus one fused broadcast Add.

  • ONNXModel gets a disabled_optimizers argument, which sets ONNX Runtime's optimization.disable_specified_optimizers session config entry.
  • FlexGlue passes {"MatMulAddFusion"} when the matcher runs on CUDA. Nothing changes on CPU or CoreML.

Can the bias be folded into the GEMM epilogue instead?

Not from the model side:

  • ONNX Runtime has no operator that fuses a plain bias into the GEMM epilogue. FusedMatMul has no bias input, and GemmFastGelu fuses only the tanh-approximated GELU.
  • The bias can't be folded into the following ops either: they are LayerNorm, the attention (after the q/k LayerNorm), and residual adds.

The separate broadcast Add that remains costs 1.36 ms per pair.

Results

Kernel profile: fp16 matcher, N = 2048, A100, ONNX Runtime 1.23 (CUPTI). GPU time: 30.42 → 28.91 ms per pair.

  • The 162 SIMT GEMMs (2.19 ms) are replaced by 162 broadcast Add kernels (1.36 ms).
  • The tensor-core GEMMs get faster with beta = 0: 5.24 → 4.85 ms.
  • fp32: 86.2 → 85.1 ms.

COLMAP matching speed:

  • Setup: exhaustive matching without geometric verification on ETH3D DSLR delivery_area + terrains + office (2132 pairs), with the keypoints of the evaluation, on an idle A100.
  • ms/pair = total matching time / pairs, including image decoding and the descriptor. Two repeats at 2k, one at 8k.
#kpts #4908 This PR FlexGlue matcher only
2048 56.0 ms 54.6 ms (−2.4%) 33.7 → 32.3 ms
8192 68.4 ms 67.2 ms (−1.8%) 45.7 → 44.4 ms

Matches: identical up to floating-point rounding; 669 instead of 668 of the 2132 pairs have matches.

Tests: feature/ and controllers/ tests pass.

ONNX Runtime's MatMulAddFusion turns each Linear layer into a Gemm, whose CUDA
kernel broadcasts the bias with an extra rank-1 GEMM on a non-tensor-core
cuBLAS kernel (162 per pair). A separate broadcast Add is faster. Adds a
disabled_optimizers argument to ONNXModel.
@sarlinpe

Copy link
Copy Markdown
Member Author

Closing this one. The gain is real but small (−1.3 ms/pair, −2.4% at 2k, −1.8% at 8k), measured only on an A100 with CUDA 12 and ONNX Runtime 1.23/1.30. The risks outweigh it:

  • The bad case is a cuBLAS kernel choice, not something inherent. ONNX Runtime's CUDA Gemm broadcasts a 1-D bias with an extra rank-1 cuBLAS GEMM. Only the kernel cuBLAS happens to pick for it on this setup (a non-tensor-core SIMT kernel) makes it slow. Both paths move about the same memory, so on other architectures or cuBLAS versions the gain may vanish.
  • Disabling the fusion would silently opt out of future fixes, for example if ONNX Runtime starts using the cuBLASLt bias epilogue for Gemm.
  • It adds a general ONNXModel knob that is only used for this runtime-specific workaround.

The finding is documented in #4908. The better fix is upstream in ONNX Runtime's CUDA Gemm: use the cuBLASLt bias epilogue, or a broadcast kernel, instead of a rank-1 GEMM.

#4915 and #4916 are re-stacked directly on #4908, and their numbers are updated.

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.

1 participant