[PR #111] [MERGED] add minimal bf16 flash attention for seqlen capability #151

New issue

Closed

opened 2026-10-26 22:20:54 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 22:20:54 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/111
Author: @rattus128
Created: 8/12/2026
Status: ✅ Merged
Merged: 8/12/2026
Merged by: @comfyanonymous

Base: main ← Head: dev/rattus/flash

📝 Commits (5)

2676f97 add minimal bf16 flash attention for seqlen capability
5b562b5 fix build issue
40cd2bf fix build issue
0867038 Fix ruff.
843eda3 Fix ruff.

📊 Changes

12 files changed (+343 additions, -7 deletions)

View changed files

📝 .github/workflows/build-wheels.yml (+7 -7)
📝 .gitmodules (+3 -0)
📝 NOTICE (+9 -0)
📝 comfy_kitchen/__init__.py (+2 -0)
📝 comfy_kitchen/backends/cuda/CMakeLists.txt (+4 -0)
📝 comfy_kitchen/backends/cuda/dlpack_bindings.cpp (+58 -0)
➕ comfy_kitchen/backends/cuda/flash_compat/ATen/cuda/CUDAGeneratorImpl.h (+12 -0)
➕ comfy_kitchen/backends/cuda/flash_compat/ATen/cuda/detail/UnpackRaw.cuh (+11 -0)
➕ comfy_kitchen/backends/cuda/ops/flash_decode.cu (+107 -0)
➕ comfy_kitchen/flash_attention.py (+75 -0)
➕ tests/test_flash_attention.py (+54 -0)
➕ third_party/flash-attention (+1 -0)

📄 Description

AR attention needs a performant seqlen capable attention kernel. Pytorch hides this option in its internals. Just use flash attention with the seqlen stuff on the API.

The seqlen variable allows cuda-graphing through the attention kernel as the CPU does not need to manage the iteration variable for sequence length increment as the AR attention iterates.

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/Comfy-Org/comfy-kitchen/pull/111 **Author:** [@rattus128](https://github.com/rattus128) **Created:** 8/12/2026 **Status:** ✅ Merged **Merged:** 8/12/2026 **Merged by:** [@comfyanonymous](https://github.com/comfyanonymous) **Base:** `main` ← **Head:** `dev/rattus/flash` --- ### 📝 Commits (5) - [`2676f97`](https://github.com/Comfy-Org/comfy-kitchen/commit/2676f971cdb2ecf10a2b68965a2bc0f00231f81c) add minimal bf16 flash attention for seqlen capability - [`5b562b5`](https://github.com/Comfy-Org/comfy-kitchen/commit/5b562b5d89baba383cd18e5208f0db7bfad52abe) fix build issue - [`40cd2bf`](https://github.com/Comfy-Org/comfy-kitchen/commit/40cd2bf3f2b6c458a8f024487929b52611af9810) fix build issue - [`0867038`](https://github.com/Comfy-Org/comfy-kitchen/commit/0867038c842c23c6acc0f2ab1092d4410b55aba7) Fix ruff. - [`843eda3`](https://github.com/Comfy-Org/comfy-kitchen/commit/843eda3e3ec1dcff15954659e36cd280d10db3cb) Fix ruff. ### 📊 Changes **12 files changed** (+343 additions, -7 deletions) <details> <summary>View changed files</summary> 📝 `.github/workflows/build-wheels.yml` (+7 -7) 📝 `.gitmodules` (+3 -0) 📝 `NOTICE` (+9 -0) 📝 `comfy_kitchen/__init__.py` (+2 -0) 📝 `comfy_kitchen/backends/cuda/CMakeLists.txt` (+4 -0) 📝 `comfy_kitchen/backends/cuda/dlpack_bindings.cpp` (+58 -0) ➕ `comfy_kitchen/backends/cuda/flash_compat/ATen/cuda/CUDAGeneratorImpl.h` (+12 -0) ➕ `comfy_kitchen/backends/cuda/flash_compat/ATen/cuda/detail/UnpackRaw.cuh` (+11 -0) ➕ `comfy_kitchen/backends/cuda/ops/flash_decode.cu` (+107 -0) ➕ `comfy_kitchen/flash_attention.py` (+75 -0) ➕ `tests/test_flash_attention.py` (+54 -0) ➕ `third_party/flash-attention` (+1 -0) </details> ### 📄 Description AR attention needs a performant seqlen capable attention kernel. Pytorch hides this option in its internals. Just use flash attention with the seqlen stuff on the API. The seqlen variable allows cuda-graphing through the attention kernel as the CPU does not need to manage the iteration variable for sequence length increment as the AR attention iterates. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

2026-10-26 22:20:54 +00:00