Skip to content

[CUDA] Add build option for TensorRT fused attention cubins - #31632

Merged
Tianlei Wu (tianleiwu) merged 3 commits into
mainfrom
tlwu/20260804/build_option_trt_cubin
Aug 5, 2026
Merged

[CUDA] Add build option for TensorRT fused attention cubins#31632
Tianlei Wu (tianleiwu) merged 3 commits into
mainfrom
tlwu/20260804/build_option_trt_cubin

Conversation

@tianleiwu

Copy link
Copy Markdown
Contributor

This change adds a build option for excluding the prebuilt TensorRT fused multi-head attention cubins from the CUDA Execution Provider. The option is enabled by default to preserve existing behavior; disabling it removes approximately 14 MB of embedded cubin data and leaves attention selection to the other available kernels and unfused fallback paths.

Key Changes

  • Add onnxruntime_USE_TRT_FUSED_ATTENTION, dependent on CUDA and enabled by default.
  • Propagate the option to CUDA provider compilation and guard TensorRT fused-attention cubin declarations, metadata, and lookup paths.
  • Filter the prebuilt TensorRT fused-attention cubin sources from both the regular CUDA provider and CUDA plugin provider when the option is disabled, while retaining the shared driver wrapper needed by sparse attention.
  • Add Windows plugin DLL version metadata and suppress build-machine RPATH embedding for packaged Linux plugin binaries.
  • Exercise onnxruntime_USE_TRT_FUSED_ATTENTION=OFF in the Windows CUDA no-cuDNN plugin build.

Testing Notes

  • git diff --check origin/main...HEAD passes.
  • The Windows CUDA no-cuDNN workflow now builds the CUDA plugin with TensorRT fused attention disabled, providing CI coverage for the opt-out configuration.
  • No local CUDA/Windows build was run in this environment.

Copilot AI 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.

Pull request overview

Adds a CUDA build-time opt-out for embedding TensorRT fused multi-head attention prebuilt cubins, reducing binary size when disabled while preserving current defaults and ensuring the plugin build remains package-friendly across platforms.

Changes:

  • Introduce onnxruntime_USE_TRT_FUSED_ATTENTION (CUDA-dependent, default ON) and propagate it as USE_TRT_FUSED_ATTENTION to the provider build.
  • Guard TensorRT fused-attention cubin declarations/metadata/lookup helpers and filter the corresponding .cubin.cc / _kernel.sm*.cc sources when the option is OFF.
  • Update CUDA plugin packaging behavior (Windows DLL version resource; Linux skip build RPATH) and add CI coverage for the opt-out configuration.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated no comments.

Show a summary per file
File Description
onnxruntime/contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/mha_runner.cu Stubs TRT fused runner support when the build option is disabled.
onnxruntime/contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/fused_multihead_attention_v2.h Guards TRT fused cubin metadata/lookup and provides null-return stubs when disabled.
onnxruntime/contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/flash_attention/fmha_flash_attention.h Guards flash-attention cubin metadata/lookup behind the build option.
onnxruntime/contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/cross_attention/fmha_cross_attention.h Guards cross-attention cubin metadata/lookup behind the build option.
cmake/onnxruntime_providers_cuda.cmake Filters out TRT fused cubin blob sources when the option is OFF.
cmake/onnxruntime_providers_cuda_plugin.cmake Applies the same filtering for plugin builds; adds Windows version resource and Linux RPATH suppression.
cmake/CMakeLists.txt Defines the new onnxruntime_USE_TRT_FUSED_ATTENTION option and emits the provider compile define when enabled.
.github/workflows/windows_cuda_no_cudnn.yml Exercises onnxruntime_USE_TRT_FUSED_ATTENTION=OFF in Windows CUDA plugin CI.

@tianleiwu
Tianlei Wu (tianleiwu) enabled auto-merge (squash) August 4, 2026 19:45
@tianleiwu
Tianlei Wu (tianleiwu) merged commit 25cc21c into main Aug 5, 2026
92 of 94 checks passed
@tianleiwu
Tianlei Wu (tianleiwu) deleted the tlwu/20260804/build_option_trt_cubin branch August 5, 2026 23:01
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.

3 participants