[PR #95] [MERGED] Add na3d / na2d neighborhood attention #142

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/95
Author: @kijai
Created: 8/6/2026
Status: ✅ Merged
Merged: 8/7/2026
Merged by: @comfyanonymous

Base: main ← Head: na

📝 Commits (6)

39c173c init natten kernels
e27faec Fixes
7a80118 Merge remote-tracking branch 'upstream/main' into na
b1fe2e1 Fixes
bebcfd8 shared validation, scalar args, sm80 gate for triton
abdd53d Update test_na.py

📊 Changes

12 files changed (+1240 additions, -0 deletions)

View changed files

📝 README.md (+2 -0)
📝 comfy_kitchen/__init__.py (+69 -0)
📝 comfy_kitchen/backends/cuda/CMakeLists.txt (+1 -0)
📝 comfy_kitchen/backends/cuda/__init__.py (+68 -0)
📝 comfy_kitchen/backends/cuda/dlpack_bindings.cpp (+39 -0)
➕ comfy_kitchen/backends/cuda/ops/na3d.cu (+433 -0)
📝 comfy_kitchen/backends/eager/__init__.py (+12 -0)
➕ comfy_kitchen/backends/eager/na.py (+179 -0)
📝 comfy_kitchen/backends/triton/__init__.py (+24 -0)
➕ comfy_kitchen/backends/triton/na.py (+155 -0)
📝 comfy_kitchen/constraints.py (+35 -0)
➕ tests/test_na.py (+223 -0)

📄 Description

Adds fused neighborhood attention (NATTEN na3d semantics, dilation 1) over
(B, T, H, W, num_heads, head_dim) tensors, plus an na2d wrapper for the singleton-T case.

backend	dtypes	head_dim	notes
cuda	fp16, bf16	16/32/48/64	fused mma kernel, sm80+
triton	fp32, fp16, bf16	≤ 128
eager	fp32, fp16, bf16	any	SDPA fallback, all devices

_{🔄 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/95 **Author:** [@kijai](https://github.com/kijai) **Created:** 8/6/2026 **Status:** ✅ Merged **Merged:** 8/7/2026 **Merged by:** [@comfyanonymous](https://github.com/comfyanonymous) **Base:** `main` ← **Head:** `na` --- ### 📝 Commits (6) - [`39c173c`](https://github.com/Comfy-Org/comfy-kitchen/commit/39c173c4aba178f4921d8d7845be5a973e401b1e) init natten kernels - [`e27faec`](https://github.com/Comfy-Org/comfy-kitchen/commit/e27faec035f23a68f833fd300d4ce734fa05391c) Fixes - [`7a80118`](https://github.com/Comfy-Org/comfy-kitchen/commit/7a80118febf2fc97947beb85e28182b66bed4864) Merge remote-tracking branch 'upstream/main' into na - [`b1fe2e1`](https://github.com/Comfy-Org/comfy-kitchen/commit/b1fe2e1a0275cfba6b6ffa22cad1780ec12b8617) Fixes - [`bebcfd8`](https://github.com/Comfy-Org/comfy-kitchen/commit/bebcfd8a345d72acd2d8b4d94ccf1509697fbecb) shared validation, scalar args, sm80 gate for triton - [`abdd53d`](https://github.com/Comfy-Org/comfy-kitchen/commit/abdd53d41754634bbb29eff9c5693099619092d4) Update test_na.py ### 📊 Changes **12 files changed** (+1240 additions, -0 deletions) <details> <summary>View changed files</summary> 📝 `README.md` (+2 -0) 📝 `comfy_kitchen/__init__.py` (+69 -0) 📝 `comfy_kitchen/backends/cuda/CMakeLists.txt` (+1 -0) 📝 `comfy_kitchen/backends/cuda/__init__.py` (+68 -0) 📝 `comfy_kitchen/backends/cuda/dlpack_bindings.cpp` (+39 -0) ➕ `comfy_kitchen/backends/cuda/ops/na3d.cu` (+433 -0) 📝 `comfy_kitchen/backends/eager/__init__.py` (+12 -0) ➕ `comfy_kitchen/backends/eager/na.py` (+179 -0) 📝 `comfy_kitchen/backends/triton/__init__.py` (+24 -0) ➕ `comfy_kitchen/backends/triton/na.py` (+155 -0) 📝 `comfy_kitchen/constraints.py` (+35 -0) ➕ `tests/test_na.py` (+223 -0) </details> ### 📄 Description Adds fused neighborhood attention (NATTEN `na3d` semantics, dilation 1) over `(B, T, H, W, num_heads, head_dim)` tensors, plus an `na2d` wrapper for the singleton-T case. | backend | dtypes | head_dim | notes | |---------|--------|----------|-------| | cuda | fp16, bf16 | 16/32/48/64 | fused mma kernel, sm80+ | | triton | fp32, fp16, bf16 | ≤ 128 | | | eager | fp32, fp16, bf16 | any | SDPA fallback, all devices | --- <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:48 +00:00