Flash-DMA is a high-performance attention implementation that integrates Flash Attention's memory efficiency with Dynamic Mask Attention's sparse computation capabilities for processing extremely long sequences in transformer models.

• Mask & Bias Support: Native support for ({1|batch_size}, {1|num_kv_heads|num_heads}, {1|query_len}, {1|key_len}) shaped attention mask and attention bias tensors
• Intelligent Computation Skipping: Block-level automatic skipping mechanism based on masks, completely bypassing computation and memory access for zero-mask blocks
• Complete Gradient Support: Built-in full gradient computation path for attention bias, supporting end-to-end training

• Dynamic Sparse Attention: Dynamically selects the most relevant keys for each query, reducing computational complexity from $O(N^2)$ to $O(N \cdot w)$ where $w \ll N$, supporting trainable sparse structures
• Memory Efficiency: Maintains Flash Attention's $O(N)$ memory complexity without instantiating the full attention matrix
• CUDA Deep Optimization: Custom CUDA kernels with shared memory aliasing, pipelined prefetching, and block skipping for high throughput and low memory access overhead
• Extremely Long Context Support: Handles 128K+ token sequences efficiently through dynamic mask windowing while preserving accuracy

We present the expected speedup of Flash-DMA over standard PyTorch SDPA under mask and bias conditions.

The following table shows the forward pass performance comparison between Flash-DMA and standard PyTorch SDPA on an NVIDIA A100-SXM4-80GB. Results are averaged over 3 runs after 2 warmup runs.

The following table shows the backward pass performance comparison between Flash-DMA and standard PyTorch SDPA on an NVIDIA A100-SXM4-80GB. Results are averaged over 3 runs after 2 warmup runs.

• Linux: Ubuntu 22.04 or later
• NVIDIA GPU: Compute Capability 8.0 or higher
• C++ Compiler: GCC 7+
• CUDA: 11.8 or later
• Python: 3.9 or later
• PyTorch: 2.5.1 or later

You can install Flash-DMA via pre-compiled wheels:

pip install flash-dmattn --no-build-isolation

Alternatively, you can compile and install from source:

git clone https://github.com/SmallDoges/flash-dmattn.git cd flash-dmattn pip install . --no-build-isolation

import torch from flash_dmattn import flash_dmattn_func_auto from flash_dmattn.utils.mask import create_mask import math # Setup batch_size, seq_len, num_heads, num_kv_heads, head_dim = 1, 256, 2, 1, 64 window_size = 128 device = torch.device('cuda') dtype = torch.bfloat16 min_dtype = torch.finfo(dtype).min # dtype minimum value # Input tensors query = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype) key = torch.randn(batch_size, seq_len, num_kv_heads, head_dim, device=device, dtype=dtype) value = torch.randn(batch_size, seq_len, num_kv_heads, head_dim, device=device, dtype=dtype) # Create bias for sparse attention attn_bias = torch.randn(batch_size, num_kv_heads, seq_len, seq_len, device=device, dtype=dtype) # Generate dynamic mask based on bias if seq_len > window_size: attn_mask = create_mask( attention_bias=attn_bias, attention_mask=None, batch_size=batch_size, query_len=seq_len, key_len=seq_len, window_size=window_size, min_dtype=min_dtype, ) # Select FDMA kernel flash_dmattn_func = flash_dmattn_func_auto(backend="cuda") # Run Flash Dynamic Mask Attention output = flash_dmattn_func( query=query, key=key, value=value, attn_mask=attn_mask, attn_bias=attn_bias, is_causal=True, softmax_scale=1.0/math.sqrt(head_dim), ) print(f"Output shape: {output.shape}") # [1, 256, 2, 64]

# Enable gradient computation query.requires_grad_(True) key.requires_grad_(True) value.requires_grad_(True) attn_bias.requires_grad_(True) # Forward pass output = flash_dmattn_func( query=query, key=key, value=value, attn_mask=attn_mask, attn_bias=attn_bias, is_causal=True, softmax_scale=1.0/math.sqrt(head_dim) ) # Backward pass loss = output.sum() loss.backward() print(f"Query gradient shape: {query.grad.shape}") print(f"Key gradient shape: {key.grad.shape}") print(f"Value gradient shape: {value.grad.shape}") print(f"Bias gradient shape: {attn_bias.grad.shape}")

Flash-DMA integrates the efficient memory access patterns of Flash Attention with the sparse computation capabilities of dynamic mask attention to achieve an efficient attention mechanism.

• 🎯 Native Mask & Bias Support: Kernels directly process ({1|batch_size}, {1|num_kv_heads|num_heads}, {1|query_len}, {1|key_len}) shaped tensors
• ⚡ Block-level Intelligent Skipping: Unified OR-reduction skipping logic based on masks, completely avoiding computation and memory access for zero blocks
• 🔄 Complete Gradient Chain: Built-in attention bias gradient computation supporting end-to-end differentiable training

1. Unified Skip Logic: Forward and backward passes use the same block-level skip decisions
2. Memory Access Optimization: K/V data loaded only when OR(mask_block) == true
3. Gradient Path Completeness: dbias gradient computation fully fused in backward kernels
4. Shared Memory Reuse: sMask ↔ sP, sBias ↔ sdS intelligent aliasing

📚 Complete documentation is available in the docs directory:

• API Reference - Complete function documentation and usage examples
• Integration Guide - Detailed technical documentation of the Flash Attention integration

# Clone with submodules git clone https://github.com/SmallDoges/flash-dmattn.git cd flash-dmattn # Build in development mode pip install -e . # Run tests to verify installation python -c "import flash_dma_cuda; print('✅ Flash DMA CUDA extension imported successfully')"

• CUDA Toolkit 11.8+
• CUTLASS library
• PyTorch with CUDA support

• SM 8.0
• SM 9.0
• SM 10.0
• SM 12.0

Note: Flash Dynamic Mask Attention requires CUDA compute capability 8.0+ for optimal performance. Earlier architectures are not supported.

Flash-DMA provides comprehensive benchmarking tools to evaluate performance across different configurations:

python benchmarks/forward_equivalence.py

Validates numerical consistency between Python reference and CUDA implementation.

python benchmarks/forward_performance.py

Compares Flash-DMA against standard SDPA across various sequence lengths and batch sizes.

python benchmarks/backward_equivalence.py

Validates numerical consistency between Python reference and CUDA implementation.

python benchmarks/backward_performance.py

Compares Flash-DMA against standard SDPA across various sequence lengths and batch sizes.

python benchmarks/grad_equivalence.py

Tests backward pass implementation and gradient equivalence.

Compilation Errors

# Ensure CUDA_HOME is set correctly echo $CUDA_HOME # Linux/Mac echo $env:CUDA_HOME # Windows PowerShell # Check CUDA toolkit version nvcc --version # Verify PyTorch CUDA support python -c "import torch; print(f'CUDA available: {torch.cuda.is_available()}')"

Import Errors

# Test basic import try: from flash_dmattn import flash_dmattn_func, get_available_backends print("✅ Flash Dynamic Mask Attention imported successfully") print(f"Available backends: {get_available_backends()}") except ImportError as e: print(f"❌ Import failed: {e}") print("Please ensure the package is properly installed with: pip install -e .")

Performance Issues

# Monitor GPU memory usage from flash_dmattn import flash_dmattn_func def print_memory_stats(): if torch.cuda.is_available(): print(f"GPU Memory: {torch.cuda.memory_allocated() / 1e9:.2f} GB") print_memory_stats() output = flash_dmattn_func(q=query, k=key, v=value, is_causal=True) print_memory_stats() # Clear cache if needed torch.cuda.empty_cache()

We welcome contributions from the community! Flash-DMA is an open-source project and we value all types of contributions.

• Report bugs: Found a bug? Please open an issue
• Request features: Have an idea for improvement? Let us know
• Submit code: Ready to contribute code? Check our Contributing Guide
• Improve docs: Help us make the documentation better

1. Fork the repository
2. Create a feature branch: git checkout -b feature-name
3. Make your changes and test them
4. Submit a pull request

For detailed instructions, see our Contributing Guide.

This project follows the Contributor Covenant Code of Conduct. By participating, you are expected to uphold this code.

This project is licensed under the BSD 3-Clause License. See LICENSE for details.

If you use Flash-DMA in your research, please cite:

@misc{shi2025trainabledynamicmasksparse, title={Trainable Dynamic Mask Sparse Attention}, author={Jingze Shi and Yifan Wu and Bingheng Wu and Yiran Peng and Liangdong Wang and Guang Liu and Yuyu Luo}, year={2025}, eprint={2508.02124}, archivePrefix={arXiv}, primaryClass={cs.AI}, url={https://arxiv.org/abs/2508.02124}, }

This project builds upon and integrates several excellent works:

• OpenSeek - Kernel development support
• Flash-Attention - Memory-efficient attention computation
• NVIDIA CUTLASS - High-performance matrix operations library

We thank the open-source community for their contributions to efficient transformer implementations. 🤗