Add backend selection and dynamic mask attention support #63
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Introduces extras_require configuration to allow users to install specific backends (triton, flex, all, dev, test) without requiring full CUDA compilation. Implements auto-detection logic that skips CUDA build when users explicitly request only Triton or Flex backends, reducing installation time and complexity for users who don't need CUDA acceleration. Maintains backward compatibility while providing more granular control over dependencies and build processes.
Initializes the flash_dmattn package with automatic backend selection between CUDA, Triton, and Flex implementations. Provides graceful fallback mechanism that prioritizes CUDA for performance, then Triton and Flex as alternatives. Includes runtime availability checks and clear error messages for missing dependencies. Enables users to explicitly specify backends or rely on automatic selection based on available installations.
Renames imported functions to use consistent naming convention with 'dmattn_func' suffix across Triton and Flex Attention implementations. Updates function call parameters to use positional arguments instead of keyword arguments for cleaner code. Removes hard exit on CUDA import failure to allow graceful degradation when some implementations are unavailable.
LoserCheems
requested review from
Evanwu1125,
SNHuan,
Copilot and
wubingheng111
and removed request for
Copilot
Expands benchmark to compare Flash Attention against CUDA, Triton, and Flex implementations of Dynamic Mask Attention. Introduces modular testing framework allowing selective benchmarking of specific implementations or head-to-head comparisons. Updates function signatures to return timing measurements directly from kernels for more accurate performance metrics. Renames variables for clarity (dt_states to zoh_states, active_mask to attn_mask) and adds comprehensive error handling for missing implementations. Includes new test configurations for window size variations and non-causal attention patterns. Provides detailed performance analysis with implementation-specific speedup calculations and overhead comparisons between different approaches.
Streamlines the detection of Triton/Flex-only installations by using simpler string matching instead of complex pattern checking. Removes unnecessary check for plain installations since the core logic focuses on whether specific extras are requested without all/dev variants.
Contributor
There was a problem hiding this comment.
Pull Request Overview
Introduce optional backend selection (Triton, Flex, CUDA) and dynamic mask attention with automatic detection, along with benchmark updates to measure performance across all backends.
- Added
should_skip_cuda_build()tosetup.pyfor auto-skipping CUDA compilation when only Triton/Flex extras are requested. - Exposed
get_available_backends()and unifiedflash_dmattn_funcdispatcher in__init__.py. - Expanded benchmarks to include Triton and Flex implementations, return execution time from functions, and parameterize test types.
Reviewed Changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated no comments.
| File | Description |
|---|---|
| setup.py | Auto-detect CUDA skip, add extras_require for Triton/Flex/all/dev/test and package data. |
| flash_dmattn/init.py | Export backend flags, helper get_available_backends(), and unified flash_dmattn_func. |
| benchmarks/benchmark_forward_performance.py | Add time/measuring support in CUDA/Triton/Flex benchmarks; parameterize test_type; update imports. |
| benchmarks/benchmark_forward_equivalence.py | Correct imports and dispatcher for Triton/Flex forward-equivalence tests. |
Comments suppressed due to low confidence (4)
benchmarks/benchmark_forward_performance.py:850
- [nitpick] The new parameters
test_type,num_runs, andwarmup_runsare not documented in the function docstring. Consider updating or adding a docstring to describe their purpose.
def run_performance_benchmark(test_type='all', num_runs=3, warmup_runs=2):
benchmarks/benchmark_forward_performance.py:188
- The
timemodule is used here but not imported at the top of the file, which will raise a NameError. Please addimport time.
start_time = time.time()
benchmarks/benchmark_forward_performance.py:641
- The
gcmodule is used here but not imported at the top of the file, leading to a NameError. Please addimport gc.
gc.collect()
benchmarks/benchmark_forward_performance.py:271
- Verify that the argument order passed to
flash_dmattn_cuda.fwdmatches the native CUDA extension signature (e.g., bias before mask or vice versa). A mismatch can silently produce incorrect outputs.
result = flash_dmattn_cuda.fwd( # type: ignore
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Introduce optional backend selection for Triton and Flex, enabling users to skip CUDA builds when not needed. Implement dynamic mask attention with automatic backend detection and graceful fallbacks. Standardize function naming for consistency across implementations while maintaining backward compatibility.