[PR #150] [MERGED] sol-attn: Full-range P quantization in the exact branch and public sol_attn_chunked #178

Closed
opened 2026-10-26 22:21:16 +00:00 by bsakaguchi · 0 comments
Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/150
Author: @kijai
Created: 9/3/2026
Status: ✅ Merged
Merged: 9/3/2026
Merged by: @comfyanonymous

Base: main ← Head: sol_exact_pquant


📝 Commits (3)

📊 Changes

18 files changed (+331 additions, -159 deletions)

View changed files

📝 comfy_kitchen/__init__.py (+6 -3)
📝 comfy_kitchen/backends/cuda/__init__.py (+6 -6)
📝 comfy_kitchen/backends/cuda/dlpack_bindings.cpp (+20 -11)
📝 comfy_kitchen/backends/cuda/sage_attention/sol_attn.cu (+13 -11)
📝 comfy_kitchen/backends/cuda/sage_attention/sol_attn_exact.cu (+30 -20)
📝 comfy_kitchen/backends/cuda/sage_attention/sol_attn_preprocess.cu (+36 -15)
📝 comfy_kitchen/backends/cuda/sage_attention/sol_attn_vtranspose.cu (+11 -7)
📝 comfy_kitchen/backends/cuda/sage_attention/sol_layout.cuh (+28 -13)
📝 comfy_kitchen/backends/eager/sol_attn.py (+1 -1)
📝 comfy_kitchen/backends/hip/__init__.py (+6 -6)
📝 comfy_kitchen/backends/hip/dlpack_bindings.cpp (+16 -9)
📝 comfy_kitchen/backends/hip/launchers.h (+1 -1)
📝 comfy_kitchen/backends/hip/sage_attention/sol_attn.hip (+14 -11)
📝 comfy_kitchen/backends/hip/sage_attention/sol_attn_exact.hip (+29 -18)
📝 comfy_kitchen/backends/hip/sage_attention/sol_attn_preprocess.hip (+30 -12)
📝 comfy_kitchen/backends/hip/sage_attention/sol_attn_vtranspose.hip (+13 -7)
📝 comfy_kitchen/backends/hip/sage_attention/sol_common.h (+16 -7)
📝 tests/test_sol_attn.py (+55 -1)

📄 Description

Small follow-up PR to the sparse attention kernels

Quantize P against the block max instead of the running max in the Sol exact branch, so that every routed block uses the full u8 range

Also exportssol_attn_chunked at top level, and adds support for fp16 attention (previously only bf16).

H3 capture, 10% keep: per-head cos 0.9830/0.9447 → 0.9853/0.9521 (mean/worst)

all-blocks relL2 vs fp32 1.96% → 1.40%. Registers unchanged, ≤3% per-call cost.


🔄 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/150 **Author:** [@kijai](https://github.com/kijai) **Created:** 9/3/2026 **Status:** ✅ Merged **Merged:** 9/3/2026 **Merged by:** [@comfyanonymous](https://github.com/comfyanonymous) **Base:** `main` ← **Head:** `sol_exact_pquant` --- ### 📝 Commits (3) - [`453c0cd`](https://github.com/Comfy-Org/comfy-kitchen/commit/453c0cdf91578cf02da72fa995f801a893d0e455) exact pquant - [`d7fe624`](https://github.com/Comfy-Org/comfy-kitchen/commit/d7fe624121428f722de000af504c7d4186af8444) fix hip test - [`e07b3b0`](https://github.com/Comfy-Org/comfy-kitchen/commit/e07b3b067335b4c82a242e2ddfd7603f70a83d0b) support fp16 attention ### 📊 Changes **18 files changed** (+331 additions, -159 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/__init__.py` (+6 -3) 📝 `comfy_kitchen/backends/cuda/__init__.py` (+6 -6) 📝 `comfy_kitchen/backends/cuda/dlpack_bindings.cpp` (+20 -11) 📝 `comfy_kitchen/backends/cuda/sage_attention/sol_attn.cu` (+13 -11) 📝 `comfy_kitchen/backends/cuda/sage_attention/sol_attn_exact.cu` (+30 -20) 📝 `comfy_kitchen/backends/cuda/sage_attention/sol_attn_preprocess.cu` (+36 -15) 📝 `comfy_kitchen/backends/cuda/sage_attention/sol_attn_vtranspose.cu` (+11 -7) 📝 `comfy_kitchen/backends/cuda/sage_attention/sol_layout.cuh` (+28 -13) 📝 `comfy_kitchen/backends/eager/sol_attn.py` (+1 -1) 📝 `comfy_kitchen/backends/hip/__init__.py` (+6 -6) 📝 `comfy_kitchen/backends/hip/dlpack_bindings.cpp` (+16 -9) 📝 `comfy_kitchen/backends/hip/launchers.h` (+1 -1) 📝 `comfy_kitchen/backends/hip/sage_attention/sol_attn.hip` (+14 -11) 📝 `comfy_kitchen/backends/hip/sage_attention/sol_attn_exact.hip` (+29 -18) 📝 `comfy_kitchen/backends/hip/sage_attention/sol_attn_preprocess.hip` (+30 -12) 📝 `comfy_kitchen/backends/hip/sage_attention/sol_attn_vtranspose.hip` (+13 -7) 📝 `comfy_kitchen/backends/hip/sage_attention/sol_common.h` (+16 -7) 📝 `tests/test_sol_attn.py` (+55 -1) </details> ### 📄 Description Small follow-up PR to the sparse attention kernels Quantize P against the block max instead of the running max in the Sol exact branch, so that every routed block uses the full u8 range Also exports`sol_attn_chunked` at top level, and adds support for fp16 attention (previously only bf16). H3 capture, 10% keep: per-head cos 0.9830/0.9447 → 0.9853/0.9521 (mean/worst) all-blocks relL2 vs fp32 1.96% → 1.40%. Registers unchanged, ≤3% per-call cost. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>
bsakaguchi 2026-10-26 22:21:16 +00:00
Sign in to join this conversation.
No labels
pull-request
No milestone
No project
No assignees
1 participant
Notifications
Due date
The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference
Comfy-Org/comfy-kitchen#178
No description provided.