[PR #266] [MERGED] Implement attention masking and deltanet for tree MTP #269

New issue

Closed

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

bsakaguchi commented

2026-10-26 22:22:17 +00:00

Owner

📋 Pull Request Information

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

Base: main ← Head: prs/adaptive-mtp

📝 Commits (6)

3e0add8 Expose non-causal prefix attention and log-sum-exp
a157d9c Rename merge internals while preserving the causal Python API
bb1f24d Use explicit row masks in speculative attention merge
ff4a65d Add tree traversal and accepted-path replay to DeltaNet
9e7e526 Test tree attention and deferred DeltaNet replay
7d2b9d6 Consolidate deferred DeltaNet on fixed 1/4/8 tree kernels

📊 Changes

14 files changed (+446 additions, -219 deletions)

View changed files

📝 comfy_kitchen/__init__.py (+2 -1)
📝 comfy_kitchen/backends/cuda/dlpack_bindings.cpp (+28 -24)
📝 comfy_kitchen/backends/cuda/ops/flash_decode.cu (+16 -14)
📝 comfy_kitchen/backends/cuda/ops/gated_delta.cu (+59 -33)
📝 comfy_kitchen/backends/hip/__init__.py (+9 -6)
📝 comfy_kitchen/backends/hip/dlpack_bindings.cpp (+31 -23)
📝 comfy_kitchen/backends/hip/launchers.h (+5 -5)
📝 comfy_kitchen/backends/hip/ops/flash_decode.hip (+14 -12)
📝 comfy_kitchen/backends/hip/ops/gated_delta.hip (+46 -22)
📝 comfy_kitchen/flash_attention.py (+46 -24)
📝 comfy_kitchen/gated_delta.py (+42 -19)
📝 comfy_kitchen/int8_decode.py (+1 -1)
📝 tests/test_flash_attention.py (+48 -31)
📝 tests/test_gated_delta_fused.py (+99 -4)

📄 Description

Tree MTP is an MTP algorithm where each draft level also provides a number of sibling backups. The verification process for this needs a causal masking scheme that can do attention for multiple candidate in parallel as separate tokens but attention/deltanet needs an appropriate masking/divergeance such that they dont attend each other.

ComfyUI PR to follow

Regression tests:

_{🔄 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/266 **Author:** [@rattus128](https://github.com/rattus128) **Created:** 10/21/2026 **Status:** ✅ Merged **Merged:** 10/21/2026 **Merged by:** [@comfyanonymous](https://github.com/comfyanonymous) **Base:** `main` ← **Head:** `prs/adaptive-mtp` --- ### 📝 Commits (6) - [`3e0add8`](https://github.com/Comfy-Org/comfy-kitchen/commit/3e0add846983ea3900eeabadeb506b34a9a7a686) Expose non-causal prefix attention and log-sum-exp - [`a157d9c`](https://github.com/Comfy-Org/comfy-kitchen/commit/a157d9cb29dc6d6b8d688ba671ac6b92fdc8a53f) Rename merge internals while preserving the causal Python API - [`bb1f24d`](https://github.com/Comfy-Org/comfy-kitchen/commit/bb1f24d62e0b84174f725e6c9086e7f4acfd8158) Use explicit row masks in speculative attention merge - [`ff4a65d`](https://github.com/Comfy-Org/comfy-kitchen/commit/ff4a65d5958cba6e2a590cc6d4f05caa9f8a892a) Add tree traversal and accepted-path replay to DeltaNet - [`9e7e526`](https://github.com/Comfy-Org/comfy-kitchen/commit/9e7e526a0897808381d0acd686da58f3c9d3e591) Test tree attention and deferred DeltaNet replay - [`7d2b9d6`](https://github.com/Comfy-Org/comfy-kitchen/commit/7d2b9d65a600992d1693c1fff416a079070d23a5) Consolidate deferred DeltaNet on fixed 1/4/8 tree kernels ### 📊 Changes **14 files changed** (+446 additions, -219 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/__init__.py` (+2 -1) 📝 `comfy_kitchen/backends/cuda/dlpack_bindings.cpp` (+28 -24) 📝 `comfy_kitchen/backends/cuda/ops/flash_decode.cu` (+16 -14) 📝 `comfy_kitchen/backends/cuda/ops/gated_delta.cu` (+59 -33) 📝 `comfy_kitchen/backends/hip/__init__.py` (+9 -6) 📝 `comfy_kitchen/backends/hip/dlpack_bindings.cpp` (+31 -23) 📝 `comfy_kitchen/backends/hip/launchers.h` (+5 -5) 📝 `comfy_kitchen/backends/hip/ops/flash_decode.hip` (+14 -12) 📝 `comfy_kitchen/backends/hip/ops/gated_delta.hip` (+46 -22) 📝 `comfy_kitchen/flash_attention.py` (+46 -24) 📝 `comfy_kitchen/gated_delta.py` (+42 -19) 📝 `comfy_kitchen/int8_decode.py` (+1 -1) 📝 `tests/test_flash_attention.py` (+48 -31) 📝 `tests/test_gated_delta_fused.py` (+99 -4) </details> ### 📄 Description Tree MTP is an MTP algorithm where each draft level also provides a number of sibling backups. The verification process for this needs a causal masking scheme that can do attention for multiple candidate in parallel as separate tokens but attention/deltanet needs an appropriate masking/divergeance such that they dont attend each other. ComfyUI PR to follow Regression tests: <img width="1236" height="444" alt="image" src="https://github.com/user-attachments/assets/cb3748f7-16c8-49a7-bfd2-f1bb7ba2d8c7" /> <img width="1236" height="370" alt="image" src="https://github.com/user-attachments/assets/30ec7c22-eed5-4c96-8c59-71966bb45d80" /> --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

2026-10-26 22:22:17 +00:00