{"id":27050538,"url":"https://github.com/naidezhujimo/triton-flashattention","last_synced_at":"2025-04-09T19:08:24.572Z","repository":{"id":286242017,"uuid":"960832579","full_name":"naidezhujimo/Triton-FlashAttention","owner":"naidezhujimo","description":"This repository contains multiple implementations of Flash Attention optimized with Triton kernels, showcasing progressive performance improvements through hardware-aware optimizations. The implementations range from basic block-wise processing to advanced techniques like FP8 quantization and prefetching","archived":false,"fork":false,"pushed_at":"2025-04-08T02:17:07.000Z","size":534,"stargazers_count":1,"open_issues_count":0,"forks_count":0,"subscribers_count":1,"default_branch":"main","last_synced_at":"2025-04-09T19:08:15.555Z","etag":null,"topics":["attention","flashattention","triton"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":null,"status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/naidezhujimo.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":null,"code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2025-04-05T06:43:20.000Z","updated_at":"2025-04-08T02:17:10.000Z","dependencies_parsed_at":"2025-04-05T07:27:55.684Z","dependency_job_id":"be3997f3-60b3-4c33-a0d4-230002dc5250","html_url":"https://github.com/naidezhujimo/Triton-FlashAttention","commit_stats":null,"previous_names":["naidezhujimo/triton-flashattention"],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/naidezhujimo%2FTriton-FlashAttention","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/naidezhujimo%2FTriton-FlashAttention/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/naidezhujimo%2FTriton-FlashAttention/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/naidezhujimo%2FTriton-FlashAttention/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/naidezhujimo","download_url":"https://codeload.github.com/naidezhujimo/Triton-FlashAttention/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":248094993,"owners_count":21046770,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["attention","flashattention","triton"],"created_at":"2025-04-05T08:17:07.373Z","updated_at":"2025-04-09T19:08:24.564Z","avatar_url":"https://github.com/naidezhujimo.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# Triton-Optimized Flash Attention Implementations\n\nThis repository contains multiple implementations of Flash Attention optimized with Triton kernels, showcasing progressive performance improvements through hardware-aware optimizations. The implementations range from basic block-wise processing to advanced techniques like FP8 quantization and prefetching.\n\n## Key Features\n- 🚀 **Three Versions of Triton-optimized Attention**:\n  - **v1**: Basic 2D grid partitioning (sequence + feature dimensions)\n  - **v2**: Multi-head support with causal masking\n  - **v3**: FP8 quantization + block prefetching\n- ⚡ **Benchmark Suite** comparing:\n  - PyTorch native implementation\n  - All Flash Attention variants\n- 📊 **Memory/Time Metrics** for:\n  - Sequence lengths up to 16,392\n  - Model dimensions up to 512\n  - Multi-head (8/16 heads) configurations\n\n## Installation\n1. Clone repository:\n```bash\ngit clone https://github.com/yourusername/triton-flash-attention.git\ncd triton-flash-attention\n```\n\n2. Install dependencies (CUDA 11.7+ required):\n```bash\npip install torch triton\n```\n\n## Usage\n\n### Basic Attention Call\n```python\nfrom Attention import call_flash_attention_v2\n\n# Input dimensions: [seq_len, num_heads, head_dim]\nq = torch.randn(1024, 8, 64, device='cuda', dtype=torch.float16)\nk = torch.randn_like(q)\nv = torch.randn_like(q)\n\noutput = call_flash_attention_v2(q, k, v, is_causal=True)\n```\n\n### Version Comparison\n| Feature               | v1          | v2          | v3          |\n|-----------------------|-------------|-------------|-------------|\n| Multi-head Support    | ❌          | ✅          | ✅          |\n| Causal Masking        | ❌          | ✅          | ✅          |\n| FP8 Quantization      | ❌          | ❌          | ✅          |\n| Block Prefetching     | ❌          | ❌          | ✅          |\n| Peak Memory (16392 seq)| 184.21MB     | 168.20MB    | 168.20MB       |\n\n### Benchmark Configuration\n```python\n# Custom benchmark setup\nconfig = {\n    'seq_len': 16384,\n    'd_model': 512,\n    'num_heads': 8,\n    'head_dim': 64\n}\nbenchmark_attention(config)\n```\n\n## Performance Results\n### Runtime Comparison (1024 sequence length)\n| Implementation        | Time (ms) | Memory (MB) |\n|-----------------------|-----------|-------------|\n| PyTorch Native        | 0.507     |    19.65   |\n| FlashAttention-v1     | 0.072     |   10.88    |\n| FlashAttention-v2     | 0.065     |   10.62    |\n| FlashAttention-v3     | 0.059     |   10.62    |\n\n### Key Optimizations\n1. **Block-wise Processing**:\n   ```python\n   # v3 block configuration\n   BLOCK_M = 128  # Query block size\n   BLOCK_N = 64   # Key/Value block size\n   ```\n2. **FP8 Quantization**:\n   ```python\n    if USE_FP8:\n        q = tl.load(\n            q_ptr + offs_m[:, None] * stride_qm + tl.arange(0, head_dim)[None, :] * stride_qh,\n            mask=(offs_m[:, None] \u003c seq_len) \u0026 (tl.arange(0, head_dim)[None, :] \u003c head_dim),\n            other=0.0\n        ).to(tl.float8e5)\n        q_scale = 127.0 / tl.max(tl.abs(q.to(tl.float32))) + 1e-6\n   ```\n3. **Prefetching**:\n   ```python\n    k_block_ptr = tl.advance(k_block_ptr, (BLOCK_N, 0))\n    v_block_ptr = tl.advance(v_block_ptr, (BLOCK_N, 0))\n   ```\n\n![1](acc1.png)\n![2](acc2.png)\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fnaidezhujimo%2Ftriton-flashattention","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fnaidezhujimo%2Ftriton-flashattention","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fnaidezhujimo%2Ftriton-flashattention/lists"}