onnxruntime
3b77775f - [CUDA] Remove TensorRT fused causal attention kernels (#29143)

Commit
22 days ago
[CUDA] Remove TensorRT fused causal attention kernels (#29143) ### Description This PR removes the TensorRT fused **causal** attention kernels (the `fmha_v2_*_Causal_*` and `fmha_v2_flash_attention_*_Causal_*` cubins) and all of the code paths that selected them from the CUDA `Attention` operator. These causal fused kernels were disabled by default (since [#14732](https://github.com/microsoft/onnxruntime/pull/14732)) and were only reachable via the opt-in `ORT_ENABLE_FUSED_CAUSAL_ATTENTION` environment variable / `TRT_CAUSAL_ATTENTION` backend bit. They used fp16 accumulation, which can cause accuracy drops, and have been superseded by flash attention, memory-efficient attention, and cuDNN SDPA. Removing them deletes ~1.27M lines of generated cubin source and simplifies the attention dispatch logic. ### Key Changes - **Removed cubin sources**: Deleted all `causal/fmha_v2_fp16_Causal_*` and `flash_attention/fmha_v2_flash_attention_fp16_Causal_*` generated cubin files (70+ files). - **Dispatch simplification** ([attention.cc](onnxruntime/contrib_ops/cuda/bert/attention.cc)): Removed the `is_unidirectional_` / causal fused-runner branch in `ComputeInternal`; the fused runner path now only handles the BERT (non-causal) case. - **Kernel options** ([attention_kernel_options.cc](onnxruntime/contrib_ops/cuda/bert/attention_kernel_options.cc), [attention_kernel_options.h](onnxruntime/contrib_ops/cuda/bert/attention_kernel_options.h)): Removed `use_trt_causal_attention_`, `UseTrtCausalAttention()`, the `TRT_CAUSAL_ATTENTION` debug print, and the `causal` argument of `SetTrtFusedKernel`. - **QKV format** ([attention_common.h](onnxruntime/contrib_ops/cpu/bert/attention_common.h), [attention_prepare_qkv.cu](onnxruntime/contrib_ops/cuda/bert/attention_prepare_qkv.cu)): Removed the `Q_K_V_BNSH_QKV_BS3NH` format and the fused-causal gemm-buffer-with-bias preparation path. - **Runner API** ([mha_runner.cu](onnxruntime/contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/mha_runner.cu), [mha_runner.h](onnxruntime/contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/mha_runner.h), [fused_multihead_attention_v2.h](onnxruntime/contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/fused_multihead_attention_v2.h)): Dropped the `causal` parameter from `FusedMHARunnerFP16v2::Create` / `IsSupported` and removed the causal kernel metadata. - **Env var removed**: `ORT_ENABLE_FUSED_CAUSAL_ATTENTION` (`kEnableFusedCausalAttention`) is no longer recognized. - **Callers updated**: [multihead_attention.cc](onnxruntime/contrib_ops/cuda/bert/multihead_attention.cc), [packed_attention.cc](onnxruntime/contrib_ops/cuda/bert/packed_attention.cc), [packed_multihead_attention.cc](onnxruntime/contrib_ops/cuda/bert/packed_multihead_attention.cc), [attention_impl.cu](onnxruntime/contrib_ops/cuda/bert/attention_impl.cu) updated to the new no-causal signatures. - **Python helpers**: Removed stale `ORT_ENABLE_FUSED_CAUSAL_ATTENTION` references from the transformers benchmark helper and stable diffusion benchmark. - **Tests updated**: [attention_op_test.cc](onnxruntime/test/contrib_ops/attention_op_test.cc) and [attention_kernel_options_test.cc](onnxruntime/test/providers/cuda/test_cases/attention_kernel_options_test.cc) no longer set/assert the causal-fused option. ### Motivation and Context The fused causal kernels were off by default, carried potential fp16-accumulation accuracy risk, and added a large amount of generated cubin source to the repo. Causal attention is already well covered by flash attention, memory-efficient attention, and cuDNN SDPA, so these kernels can be safely removed to reduce binary size and simplify maintenance. ### Testing - Build the CUDA EP and run the attention contrib op tests (`ContribOpAttentionTest.*`, including `Causal_EmptyPastState`). - Run `AttentionKernelOptionsTest.*` to verify the kernel-option parsing no longer references the causal backend.
Author
Parents
Loading