[PR #42] [CLOSED] Add Sage Attention (CORE-74) #99

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/42
Author: @contentis
Created: 5/20/2026
Status: ❌ Closed

Base: main ← Head: sage_attn

📝 Commits (10+)

2bddbf9 poc sage attention
1d1cdef Enhance CUDA backend for SageAttention: update CMake configurations for new FP8 kernels, improve error handling in sage_sdpa, and streamline quantization kernels. Also, fix exception types in tests for better clarity.
926f719 formatting
1285876 clean up and fixes
9883175 build fix
aa4edc7 speed
23ca166 acc fix
e8066d5 ruff format
55f594e add f32 input for sage + readme update
d0f5a79 Merge upstream/main into sage_attn

📊 Changes

21 files changed (+5387 additions, -11 deletions)

View changed files

📝 NOTICE (+12 -0)
📝 README.md (+34 -0)
📝 comfy_kitchen/__init__.py (+32 -0)
📝 comfy_kitchen/backends/cuda/CMakeLists.txt (+21 -0)
📝 comfy_kitchen/backends/cuda/dlpack_bindings.cpp (+299 -0)
📝 comfy_kitchen/backends/cuda/float_utils.cuh (+36 -11)
➕ comfy_kitchen/backends/cuda/sage_attention/attn_utils.cuh (+1013 -0)
➕ comfy_kitchen/backends/cuda/sage_attention/cp_async.cuh (+150 -0)
➕ comfy_kitchen/backends/cuda/sage_attention/math.cuh (+170 -0)
➕ comfy_kitchen/backends/cuda/sage_attention/mma.cuh (+660 -0)
➕ comfy_kitchen/backends/cuda/sage_attention/numeric_conversion.cuh (+159 -0)
➕ comfy_kitchen/backends/cuda/sage_attention/permuted_smem.cuh (+213 -0)
➕ comfy_kitchen/backends/cuda/sage_attention/qk_int_sv_f8_cuda_sm89.cuh (+821 -0)
➕ comfy_kitchen/backends/cuda/sage_attention/quant_qk_int8.cu (+552 -0)
➕ comfy_kitchen/backends/cuda/sage_attention/quant_v_fp8.cu (+191 -0)
➕ comfy_kitchen/backends/cuda/sage_attention/sage_attn_launcher.cu (+111 -0)
➕ comfy_kitchen/sage_attention.py (+213 -0)
📝 tests/__init__.py (+1 -0)
➕ tests/sage_triton_ref/__init__.py (+2 -0)
➕ tests/sage_triton_ref/quant_per_thread.py (+221 -0)

...and 1 more files

📄 Description

This PR adds the kernel for SageAttention 2 using FP8 and slighlty updated quantization kernels for improved end-to-end performance.

_{🔄 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/42 **Author:** [@contentis](https://github.com/contentis) **Created:** 5/20/2026 **Status:** ❌ Closed **Base:** `main` ← **Head:** `sage_attn` --- ### 📝 Commits (10+) - [`2bddbf9`](https://github.com/Comfy-Org/comfy-kitchen/commit/2bddbf910d8b5590a8f75d6773cf6aeb54ba2535) poc sage attention - [`1d1cdef`](https://github.com/Comfy-Org/comfy-kitchen/commit/1d1cdef5e6c3f9c21dcf3bc6dac2206f3c4e4c47) Enhance CUDA backend for SageAttention: update CMake configurations for new FP8 kernels, improve error handling in sage_sdpa, and streamline quantization kernels. Also, fix exception types in tests for better clarity. - [`926f719`](https://github.com/Comfy-Org/comfy-kitchen/commit/926f719c2c26c96a59e58dd8b94a82791b1242a5) formatting - [`1285876`](https://github.com/Comfy-Org/comfy-kitchen/commit/12858765d1387c2a3b82e1b96a687650b789d518) clean up and fixes - [`9883175`](https://github.com/Comfy-Org/comfy-kitchen/commit/9883175017bbb33b354bdf046e287f6171db1459) build fix - [`aa4edc7`](https://github.com/Comfy-Org/comfy-kitchen/commit/aa4edc7791734b77c785ea99fc81771215c59394) speed - [`23ca166`](https://github.com/Comfy-Org/comfy-kitchen/commit/23ca1660035d86c9a73094d5b34751a2b5a95f79) acc fix - [`e8066d5`](https://github.com/Comfy-Org/comfy-kitchen/commit/e8066d5170db34071553d05794cd07ed8969e27c) ruff format - [`55f594e`](https://github.com/Comfy-Org/comfy-kitchen/commit/55f594ea5c264a59184bbf9ef3dc3c3f09347005) add f32 input for sage + readme update - [`d0f5a79`](https://github.com/Comfy-Org/comfy-kitchen/commit/d0f5a793bf02d38bb03711c2e4a439fbcb05525a) Merge upstream/main into sage_attn ### 📊 Changes **21 files changed** (+5387 additions, -11 deletions) <details> <summary>View changed files</summary> 📝 `NOTICE` (+12 -0) 📝 `README.md` (+34 -0) 📝 `comfy_kitchen/__init__.py` (+32 -0) 📝 `comfy_kitchen/backends/cuda/CMakeLists.txt` (+21 -0) 📝 `comfy_kitchen/backends/cuda/dlpack_bindings.cpp` (+299 -0) 📝 `comfy_kitchen/backends/cuda/float_utils.cuh` (+36 -11) ➕ `comfy_kitchen/backends/cuda/sage_attention/attn_utils.cuh` (+1013 -0) ➕ `comfy_kitchen/backends/cuda/sage_attention/cp_async.cuh` (+150 -0) ➕ `comfy_kitchen/backends/cuda/sage_attention/math.cuh` (+170 -0) ➕ `comfy_kitchen/backends/cuda/sage_attention/mma.cuh` (+660 -0) ➕ `comfy_kitchen/backends/cuda/sage_attention/numeric_conversion.cuh` (+159 -0) ➕ `comfy_kitchen/backends/cuda/sage_attention/permuted_smem.cuh` (+213 -0) ➕ `comfy_kitchen/backends/cuda/sage_attention/qk_int_sv_f8_cuda_sm89.cuh` (+821 -0) ➕ `comfy_kitchen/backends/cuda/sage_attention/quant_qk_int8.cu` (+552 -0) ➕ `comfy_kitchen/backends/cuda/sage_attention/quant_v_fp8.cu` (+191 -0) ➕ `comfy_kitchen/backends/cuda/sage_attention/sage_attn_launcher.cu` (+111 -0) ➕ `comfy_kitchen/sage_attention.py` (+213 -0) 📝 `tests/__init__.py` (+1 -0) ➕ `tests/sage_triton_ref/__init__.py` (+2 -0) ➕ `tests/sage_triton_ref/quant_per_thread.py` (+221 -0) _...and 1 more files_ </details> ### 📄 Description This PR adds the kernel for SageAttention 2 using FP8 and slighlty updated quantization kernels for improved end-to-end performance. <img width="2974" height="1910" alt="image" src="https://github.com/user-attachments/assets/183fb5de-8246-4333-b132-f54d9e3bf807" /> --- <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:23 +00:00