onnxruntime/onnxruntime/core/framework
Vincent Wang b7408f7389
[ORTModule] ATen Efficient Attention and Triton Flash Attention (#17959)
This PR is to support efficient attention and flash attention in
ORTModule, including:
- Use ATen to call efficient attention, which requires PyTorch 2.2.0 dev
or newer. ORTMODULE_USE_EFFICIENT_ATTENTION=1 to enable.
- Integrate Triton Flash attention, which requires
triton==2.0.0.dev20221202. Need A100 or H100.
ORTMODULE_USE_FLASH_ATTENTION=1 to enable.
- A python transformer tool to match sub-graph by config and write
transformer quickly.

Current transformers supports attention mask for both efficient attn and
flash attn, and dropout for efficient attn only. To support more
training scenarios (such as causal mask in GPT2), more transformers need
to be added.

The feature is guarded by system environment variables, it won't effect
any current behavior if not enabled. Since it requires specific
PyTorch/Triton versions, related tests is not added for now.
2023-10-27 10:29:27 +08:00
..
allocation_planner.cc Support inplace update for PythonOp/Grad (#17687) 2023-10-10 21:36:45 -07:00
allocation_planner.h ExecutionProvider API refactor - replace OrtMemoryInfo with OrtDevice (#15618) 2023-05-01 10:06:00 -07:00
allocator.cc OrtMemoryInfo: support new name "WebGPU_Buffer" (#17469) 2023-09-08 16:37:35 -07:00
allocator_stats.h
allocator_utils.cc remove AllocatorMgr class (#16509) 2023-06-28 15:43:19 -07:00
allocator_utils.h remove AllocatorMgr class (#16509) 2023-06-28 15:43:19 -07:00
arena_extend_strategy.h
bfc_arena.cc New configuration to limit the arena extension (#15983) 2023-05-25 02:19:07 -07:00
bfc_arena.h New configuration to limit the arena extension (#15983) 2023-05-25 02:19:07 -07:00
callback.cc
callback.h
compute_capability.h Run clang-format in CI (#15524) 2023-04-18 09:26:58 -07:00
config_options.cc
config_options.h
copy.cc
copy.h
customregistry.cc
data_transfer.cc
data_transfer.h
data_transfer_manager.cc Run clang-format in CI (#15524) 2023-04-18 09:26:58 -07:00
data_transfer_manager.h
data_transfer_utils.h
data_types.cc [C#, CPP] Introduce Float16/BFloat16 support and tests for C#, C++ (#16506) 2023-07-14 10:46:52 -07:00
debug_node_inputs_outputs_utils.cc dump tensor statistics (#15761) 2023-08-16 10:53:48 -07:00
debug_node_inputs_outputs_utils.h dump tensor statistics (#15761) 2023-08-16 10:53:48 -07:00
device_stream_collection.cc clean unused parameter in ORT_UNUSED_PARAMETER (#16538) 2023-07-07 13:20:36 -07:00
device_stream_collection.h ExecutionProvider API refactor - move allocator from EP level to SessionState level and indexed by OrtDevice (#15833) 2023-06-19 17:44:45 -07:00
element_type_lists.h Introduce float 8 types (#14731) 2023-05-30 13:25:58 -07:00
empty.cc
endian_utils.cc
endian_utils.h
error_code.cc Run clang-format in CI (#15524) 2023-04-18 09:26:58 -07:00
error_code_helper.h Run clang-format in CI (#15524) 2023-04-18 09:26:58 -07:00
ex_lib_loader.cc
ex_lib_loader.h
execution_frame.cc fix memory profile build (#16177) 2023-05-31 16:08:14 -07:00
execution_frame.h ExecutionProvider API refactor - replace OrtMemoryInfo with OrtDevice (#15618) 2023-05-01 10:06:00 -07:00
execution_plan_base.h ExecutionProvider API refactor - replace OrtMemoryInfo with OrtDevice (#15618) 2023-05-01 10:06:00 -07:00
execution_provider.cc ExecutionProvider API refactor - move allocator from EP level to SessionState level and indexed by OrtDevice (#15833) 2023-06-19 17:44:45 -07:00
execution_providers.h ExecutionProvider API refactor - move allocator from EP level to SessionState level and indexed by OrtDevice (#15833) 2023-06-19 17:44:45 -07:00
execution_steps.cc Run clang-format in CI (#15524) 2023-04-18 09:26:58 -07:00
execution_steps.h Run clang-format in CI (#15524) 2023-04-18 09:26:58 -07:00
fallback_cpu_capability.cc [ORTModule] ATen Efficient Attention and Triton Flash Attention (#17959) 2023-10-27 10:29:27 +08:00
fallback_cpu_capability.h Run clang-format in CI (#15524) 2023-04-18 09:26:58 -07:00
feeds_fetches_manager.cc
feeds_fetches_manager.h Run clang-format in CI (#15524) 2023-04-18 09:26:58 -07:00
func_kernel.cc
func_kernel.h Run clang-format in CI (#15524) 2023-04-18 09:26:58 -07:00
fuse_nodes_funcs.cc Run clang-format in CI (#15524) 2023-04-18 09:26:58 -07:00
fuse_nodes_funcs.h Run clang-format in CI (#15524) 2023-04-18 09:26:58 -07:00
graph_partitioner.cc Functions Ahead Of Time inlininng (#17764) 2023-10-23 17:42:20 -07:00
graph_partitioner.h Functions Ahead Of Time inlininng (#17764) 2023-10-23 17:42:20 -07:00
iexecutor.h ExecutionProvider API refactor - replace OrtMemoryInfo with OrtDevice (#15618) 2023-05-01 10:06:00 -07:00
kernel_def_builder.cc Run clang-format in CI (#15524) 2023-04-18 09:26:58 -07:00
kernel_lookup.h
kernel_registry.cc Refine cast optimizer for safety (#15658) 2023-04-28 09:32:54 -07:00
kernel_registry_manager.cc Resize and EP specific transpose optimization updates (#17664) 2023-09-29 08:11:36 +10:00
kernel_registry_manager.h
kernel_type_str_resolver.cc Introduce float 8 types (#14731) 2023-05-30 13:25:58 -07:00
kernel_type_str_resolver.h Fix compilation with newer flatbuffers (#17164) 2023-08-29 10:28:26 -07:00
kernel_type_str_resolver_utils.cc Re-organize the transpose optimization and layout transformation files. (#16246) 2023-07-07 08:24:47 +10:00
kernel_type_str_resolver_utils.h
library_handles.cc
library_handles.h Run clang-format in CI (#15524) 2023-04-18 09:26:58 -07:00
math.h
mem_buffer.h
mem_pattern.h ExecutionProvider API refactor - replace OrtMemoryInfo with OrtDevice (#15618) 2023-05-01 10:06:00 -07:00
mem_pattern_planner.h Run clang-format in CI (#15524) 2023-04-18 09:26:58 -07:00
memcpy.cc
memcpy.h
memory_info.cc fix memory profile build (#16177) 2023-05-31 16:08:14 -07:00
memory_info.h fix memory profile build (#16177) 2023-05-31 16:08:14 -07:00
mldata_type_utils.cc
mldata_type_utils.h
murmurhash3.cc
murmurhash3.h
node_index_info.cc Run clang-format in CI (#15524) 2023-04-18 09:26:58 -07:00
node_index_info.h
onnxruntime_map_type_info.cc Introduce float 8 types (#14731) 2023-05-30 13:25:58 -07:00
onnxruntime_map_type_info.h Run clang-format in CI (#15524) 2023-04-18 09:26:58 -07:00
onnxruntime_optional_type_info.cc Implement Optional Metadata support and C# test support (#15314) 2023-04-11 09:41:59 -07:00
onnxruntime_optional_type_info.h Run clang-format in CI (#15524) 2023-04-18 09:26:58 -07:00
onnxruntime_sequence_type_info.cc Implement Optional Metadata support and C# test support (#15314) 2023-04-11 09:41:59 -07:00
onnxruntime_sequence_type_info.h Run clang-format in CI (#15524) 2023-04-18 09:26:58 -07:00
onnxruntime_typeinfo.cc Introduce float 8 types (#14731) 2023-05-30 13:25:58 -07:00
onnxruntime_typeinfo.h Run clang-format in CI (#15524) 2023-04-18 09:26:58 -07:00
op_kernel.cc ExecutionProvider API refactor - replace OrtMemoryInfo with OrtDevice (#15618) 2023-05-01 10:06:00 -07:00
op_kernel_context_internal.h
op_kernel_info.cc ExecutionProvider API refactor - move allocator from EP level to SessionState level and indexed by OrtDevice (#15833) 2023-06-19 17:44:45 -07:00
op_kernel_type_control_utils.h Run clang-format in CI (#15524) 2023-04-18 09:26:58 -07:00
op_node_proto_helper.cc
ort_stl_allocator.h
ort_value_name_idx_map.h
ort_value_pattern_planner.cc Run clang-format in CI (#15524) 2023-04-18 09:26:58 -07:00
ort_value_pattern_planner.h ExecutionProvider API refactor - replace OrtMemoryInfo with OrtDevice (#15618) 2023-05-01 10:06:00 -07:00
ort_value_tensor_slicer.cc
ort_value_tensor_slicer.h
partial_graph_execution_state.cc Enable conditional optimization automatically (#15885) 2023-05-23 13:08:05 +08:00
partial_graph_execution_state.h
prepacked_weights.cc
prepacked_weights.h Refactor prepack buffer code (#16280) 2023-06-08 14:42:02 -07:00
prepacked_weights_container.cc remove AllocatorMgr class (#16509) 2023-06-28 15:43:19 -07:00
prepacked_weights_container.h
print_tensor_statistics_utils.h dump tensor statistics (#15761) 2023-08-16 10:53:48 -07:00
print_tensor_utils.h Enable -Wshorten-64-to-32 warning if available. (#16524) 2023-07-07 08:11:44 -07:00
program_region.h
random_generator.cc
random_generator.h
random_seed.cc Run clang-format in CI (#15524) 2023-04-18 09:26:58 -07:00
random_seed.h Run clang-format in CI (#15524) 2023-04-18 09:26:58 -07:00
run_options.cc
sequential_execution_plan.h ExecutionProvider API refactor - replace OrtMemoryInfo with OrtDevice (#15618) 2023-05-01 10:06:00 -07:00
sequential_executor.cc Run clang-format in CI (#15524) 2023-04-18 09:26:58 -07:00
sequential_executor.h Fix memory profiler (#14695) 2023-02-23 18:05:21 +08:00
session_options.cc
session_options.h Add support for specifying a custom logging function per session. (#17727) 2023-09-29 19:46:55 -07:00
session_state.cc Add prepacked weights container to subgraphs (#17671) 2023-09-26 12:01:41 -07:00
session_state.h Fix compilation with newer flatbuffers (#17164) 2023-08-29 10:28:26 -07:00
session_state_utils.cc ORT_ENFORCE on the iterator must come before iterator is dereferenced. (#17265) 2023-08-23 17:20:01 -07:00
session_state_utils.h Run clang-format in CI (#15524) 2023-04-18 09:26:58 -07:00
simple_tensor_allocator.cc ExecutionProvider API refactor - replace OrtMemoryInfo with OrtDevice (#15618) 2023-05-01 10:06:00 -07:00
simple_tensor_allocator.h ExecutionProvider API refactor - replace OrtMemoryInfo with OrtDevice (#15618) 2023-05-01 10:06:00 -07:00
sparse_tensor.cc
sparse_utils.cc Run clang-format in CI (#15524) 2023-04-18 09:26:58 -07:00
sparse_utils.h Run clang-format in CI (#15524) 2023-04-18 09:26:58 -07:00
stream_execution_context.cc create a stream in DeviceStreamCollection for memory pattern (#15426) 2023-04-17 10:06:55 -07:00
stream_execution_context.h Enable -Wshorten-64-to-32 warning if available. (#16524) 2023-07-07 08:11:44 -07:00
tensor.cc Handle string initializers in constant folding (#17422) 2023-09-27 21:15:58 +10:00
tensor_allocator.cc ExecutionProvider API refactor - replace OrtMemoryInfo with OrtDevice (#15618) 2023-05-01 10:06:00 -07:00
tensor_allocator.h ExecutionProvider API refactor - replace OrtMemoryInfo with OrtDevice (#15618) 2023-05-01 10:06:00 -07:00
tensor_allocator_with_mem_pattern.h ExecutionProvider API refactor - replace OrtMemoryInfo with OrtDevice (#15618) 2023-05-01 10:06:00 -07:00
tensor_external_data_info.cc
tensor_external_data_info.h
tensor_shape.cc Run clang-format in CI (#15524) 2023-04-18 09:26:58 -07:00
tensor_type_and_shape.cc Custom op shape inference API (#17737) 2023-10-13 12:57:42 -07:00
tensor_type_and_shape.h Run clang-format in CI (#15524) 2023-04-18 09:26:58 -07:00
tensorprotoutils.cc Handle string initializers in constant folding (#17422) 2023-09-27 21:15:58 +10:00
tensorprotoutils.h Handle string initializers in constant folding (#17422) 2023-09-27 21:15:58 +10:00
TensorSeq.h Enable Opset11 Sequence Ops on DirectML, and make the CPU implementations agnostic to backend EP (#14442) 2023-02-21 18:08:28 -08:00
transform_layout_functions.h Re-organize the transpose optimization and layout transformation files. (#16246) 2023-07-07 08:24:47 +10:00
transpose_helper.cc Parallelize Transpose (#16854) 2023-08-07 14:25:53 -07:00
transpose_helper.h Parallelize Transpose (#16854) 2023-08-07 14:25:53 -07:00
tunable.h Update ck and enable test (#16383) 2023-08-22 11:08:55 +08:00
tuning_context.h remove gsl to make status.h independent from gsl (#17402) 2023-09-13 21:47:43 -07:00
tuning_context_impl.h Run clang-format in CI (#15524) 2023-04-18 09:26:58 -07:00
tuning_results.h Add TuningContext for TunableOp (#14557) 2023-02-10 14:27:43 +08:00
utils.cc [ORTModule] ATen Efficient Attention and Triton Flash Attention (#17959) 2023-10-27 10:29:27 +08:00
utils.h [ORTModule] ATen Efficient Attention and Triton Flash Attention (#17959) 2023-10-27 10:29:27 +08:00