Comfy-Org/comfy-kitchen

Fork 0

mirror of https://github.com/Comfy-Org/comfy-kitchen.git synced 2026-11-04 23:37:42 +00:00

[PR #179] Add draft_attention for block-sparse video attention #196

New issue

Open

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

bsakaguchi commented

2026-10-26 22:21:29 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/179
Author: @Yusux
Created: 9/16/2026
Status: 🔄 Open

Base: main ← Head: anemoi-sparse-attention

📝 Commits (6)

55c125a feat(attention): add anemoi block-sparse attention
529e46c fix(kernels): launch the Ada set on every served SM
6588350 feat(attention): add native SM120 D64 support
2817e13 fix(kernels): CUDA 12.8 compatibility for SM120 instantiations
6589e66 Merge fork/main into anemoi-sparse-attention
b489f25 refactor: Rename Anemoi implementation to Draft

📊 Changes

93 files changed (+24455 additions, -2 deletions)

View changed files

📝 README.md (+1 -0)
📝 comfy_kitchen/__init__.py (+74 -0)
📝 comfy_kitchen/backends/cuda/CMakeLists.txt (+12 -2)
📝 comfy_kitchen/backends/cuda/__init__.py (+125 -0)
📝 comfy_kitchen/backends/cuda/dlpack_bindings.cpp (+7 -0)
➕ comfy_kitchen/backends/cuda/draft/CMakeLists.txt (+131 -0)
➕ comfy_kitchen/backends/cuda/draft/assembly/output_assembly.cu (+756 -0)
➕ comfy_kitchen/backends/cuda/draft/attention/ada/LICENSE (+201 -0)
➕ comfy_kitchen/backends/cuda/draft/attention/ada/attention_decl.cuh (+33 -0)
➕ comfy_kitchen/backends/cuda/draft/attention/ada/instantiations/inst_k64_d128.cu (+31 -0)
➕ comfy_kitchen/backends/cuda/draft/attention/ada/instantiations/inst_k64_d128_int8_dense.cu (+16 -0)
➕ comfy_kitchen/backends/cuda/draft/attention/ada/instantiations/inst_k64_d64.cu (+31 -0)
➕ comfy_kitchen/backends/cuda/draft/attention/ada/instantiations/inst_k64_d64_int8_dense.cu (+16 -0)
➕ comfy_kitchen/backends/cuda/draft/attention/ada/instantiations/inst_q128_k64_d128.cu (+19 -0)
➕ comfy_kitchen/backends/cuda/draft/attention/ada/instantiations/inst_q128_k64_d128_fp16.cu (+13 -0)
➕ comfy_kitchen/backends/cuda/draft/attention/ada/instantiations/inst_q128_k64_d128_int8.cu (+13 -0)
➕ comfy_kitchen/backends/cuda/draft/attention/ada/instantiations/inst_q128_k64_d64.cu (+19 -0)
➕ comfy_kitchen/backends/cuda/draft/attention/ada/instantiations/inst_q128_k64_d64_fp16.cu (+13 -0)
➕ comfy_kitchen/backends/cuda/draft/attention/ada/instantiations/inst_q128_k64_d64_int8.cu (+13 -0)
➕ comfy_kitchen/backends/cuda/draft/attention/ada/k64_attention_decl.cuh (+98 -0)

...and 73 more files

📄 Description

Summary

Adds comfy_kitchen.draft_attention, a DraftMap-based block-sparse self-attention operation for packed video sequences, with a full-precision eager reference and native CUDA executors for NVIDIA GPUs from SM89 onward.

The implementation is adapted from anemoi-project/anemoi revision bcb18202a4afd006749700370634de4819f9cc19. It consolidates option validation, layout planning, backend dispatch, native availability checks, and the CUDA executor behind one public Comfy Kitchen operation.

What changed

Adds the public APIs:
- comfy_kitchen.draft_attention(...)
- comfy_kitchen.draft_attention_is_available(device)
Registers eager and CUDA implementations with capability-aware dispatch and constraint metadata.
Adds a device-portable full-precision eager reference for CPU and unsupported accelerators/devices.
Adds shared validation and packed video-layout construction without persistent GPU tensor caches.
Adds the migrated native CUDA plan, routing, draft, preparation, attention, and output-assembly stages.
Serves capabilities [89, 120) with the Ada kernel set; SM120 uses the SM120a executor.
Adds build-system integration and backend capability reporting.
Adds tests/test_draft_attention.py covering the public contract, option and layout validation, dispatch, availability, eager correctness, native stages, and CUDA reference execution.

Public behavior

Inputs are equal-shaped BTHD fp16/bf16 tensors with batch size one plus a packed video shape and prefix length. The default profile is Q64 with 80% sparsity and an INT8 middle phase.

The native path is selected only when the complete executor was compiled for the current CUDA device and the call satisfies its constraints. CPU, unsupported accelerators, and CUDA devices without a matching compiled executor fall through to the full-precision eager reference instead of failing dispatch.

Benchmarks

End-to-end protocol

MiniMax H3 video generation through ComfyUI, with Yusux/ComfyUI@e65ddf0 and Yusux/comfy-kitchen-anemoi@2817e13 using:

minimax_h3_fl2va_pruned_fp8_scaled.safetensors (UNet)
qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors (text encoder)
minimax_h3_video_vae_fp16.safetensors + minimax_h3_audio_vae_fp32.safetensors (VAEs)

These measurements were collected before the Anemoi-to-Draft rename. The tables use the current Draft name for the same implementation; E2E benchmarks were not rerun after the naming-only changes.

All runs use 1344×768, 124 frames, 24 FPS, and 50 sampling steps, with the res_multistep sampler and simple scheduler. The three scenes are city (seed 0), ocean (seed 0), and forest (seed 1234), with identical prompts across systems. Each scene/configuration has one six-step warmup followed by one 50-step formal run on a persistent server; warmup time is excluded. Reported E2E wall time runs from prompt submission to completion and includes sampling, video/audio VAE decode, and video saving. Completion is polled every two seconds.

All sparse profiles use Q64 and start_percent=0.2. This setting is converted to a sigma threshold rather than an integer step count. In all 27 formal sparse runs, Draft attention starts after the initial 11-step dense phase. Drop denotes the configured sparsity ratio; NVFP4/INT8 percentages describe the precision split of retained sparse work, not the model weights. The attention video grid is (37, 24, 42), with 56 heads, head dimension 128, and 473–479 prefix tokens (depending on the prompt).

Cross-architecture E2E: Q64, Drop80, 100% INT8

Means across the three scenes (one formal run per scene/configuration). E2E speedups below are Dense mean time divided by sparse mean time:

GPU	Architecture	PyTorch	Dense E2E (s)	Draft E2E (s)	Speedup vs Dense	RGB PSNR vs Dense (dB)
RTX 4090 (48 GB)	SM89	2.12.1+cu130	1037.20	615.58	1.685×	20.46
H200	SM90	2.12.1+cu130	392.58	236.43	1.660×	20.51
RTX 5090	SM120	2.11.0+cu130	704.84	388.42	1.815×	20.88

SM89 and SM90 use the native Ada-family INT8 executor; the mixed NVFP4 attention profiles below were tested only on SM120. PSNR is computed over decoded RGB frames against each scene's Dense output on the same GPU, then averaged across scenes. It measures output similarity, not perceptual quality or cross-GPU numerical equivalence; audio is not scored.

RTX 5090 E2E: seven Q64 sparse profiles

All seven profiles use INT8 prefix K/V and prefix queries.

Drop	Retained NVFP4 / INT8	Mean E2E (s)	Speedup vs Dense	RGB PSNR vs Dense (dB)
—	Dense reference	704.84	1.000×	—
80%	0% / 100%	388.42	1.815×	20.88
60%	60% / 40%	413.25	1.706×	21.22
60%	75% / 25%	410.46	1.717×	20.60
70%	60% / 40%	399.79	1.763×	20.81
70%	75% / 25%	397.17	1.775×	19.83
80%	60% / 40%	386.47	1.824×	20.04
80%	75% / 25%	383.10	1.840×	19.44

At matched Drop80, NV75/I25 reduces mean wall time by 5.32 s versus pure INT8 (about 1.014× additional speedup). Drop60/70 retain more sparse work and remain slower than Drop80 INT8.

RTX 5090 Q64 preparation / mainloop breakdown

Separate operator profiling at the city attention geometry: B=1, T=37775, H=56, D=128, video grid (37, 24, 42), prefix length 479. Values are means of three CUDA-profiler captures per configuration. These are kernel-duration sums for preparation and native attention mainloop only, not full operator wall time and not the 50-step E2E measurement. Other stages and host overhead are excluded; sums below are rounded from unrounded values.

Drop	Retained NVFP4 / INT8	Preparation (ms)	Mainloop (ms)	Prep + mainloop (ms)
80%	0% / 100%	3.393	15.046	18.439
60%	60% / 40%	4.479	24.583	29.062
60%	75% / 25%	4.495	23.369	27.864
70%	60% / 40%	4.463	18.863	23.325
70%	75% / 25%	4.418	17.698	22.115
80%	60% / 40%	4.395	12.881	17.276
80%	75% / 25%	4.413	12.377	16.790

Mixed NVFP4/INT8 adds approximately 1 ms of preparation, but at matched Drop80 its mainloop saving offsets that cost: prep + mainloop falls by 6.31% for NV60/I40 and 8.94% for NV75/I25. This component-level saving is diluted in E2E by the dense-first sampling portion and work outside sparse attention.

Provenance and licensing

The migrated implementation is adapted from the Apache-2.0 Anemoi source at revision bcb18202a4afd006749700370634de4819f9cc19. SPDX headers, Ada/Blackwell license files, Blackwell provenance, and per-file upstream attribution are preserved.

_{🔄 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/179 **Author:** [@Yusux](https://github.com/Yusux) **Created:** 9/16/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `anemoi-sparse-attention` --- ### 📝 Commits (6) - [`55c125a`](https://github.com/Comfy-Org/comfy-kitchen/commit/55c125aaf6f6756a8c316540788384a951cc4054) feat(attention): add anemoi block-sparse attention - [`529e46c`](https://github.com/Comfy-Org/comfy-kitchen/commit/529e46cfb7052287d0dee1a251ca34c21f4a597a) fix(kernels): launch the Ada set on every served SM - [`6588350`](https://github.com/Comfy-Org/comfy-kitchen/commit/658835013c5487eda01a16a1fb43d815645d33b5) feat(attention): add native SM120 D64 support - [`2817e13`](https://github.com/Comfy-Org/comfy-kitchen/commit/2817e131b4cdf45710a374d2f1d14d836a73a772) fix(kernels): CUDA 12.8 compatibility for SM120 instantiations - [`6589e66`](https://github.com/Comfy-Org/comfy-kitchen/commit/6589e66f3fbfacdac166dd37a393c3b6ac0b5fe1) Merge fork/main into anemoi-sparse-attention - [`b489f25`](https://github.com/Comfy-Org/comfy-kitchen/commit/b489f253678c3e2adbd1c2f51e31a28098f961b5) refactor: Rename Anemoi implementation to Draft ### 📊 Changes **93 files changed** (+24455 additions, -2 deletions) <details> <summary>View changed files</summary> 📝 `README.md` (+1 -0) 📝 `comfy_kitchen/__init__.py` (+74 -0) 📝 `comfy_kitchen/backends/cuda/CMakeLists.txt` (+12 -2) 📝 `comfy_kitchen/backends/cuda/__init__.py` (+125 -0) 📝 `comfy_kitchen/backends/cuda/dlpack_bindings.cpp` (+7 -0) ➕ `comfy_kitchen/backends/cuda/draft/CMakeLists.txt` (+131 -0) ➕ `comfy_kitchen/backends/cuda/draft/assembly/output_assembly.cu` (+756 -0) ➕ `comfy_kitchen/backends/cuda/draft/attention/ada/LICENSE` (+201 -0) ➕ `comfy_kitchen/backends/cuda/draft/attention/ada/attention_decl.cuh` (+33 -0) ➕ `comfy_kitchen/backends/cuda/draft/attention/ada/instantiations/inst_k64_d128.cu` (+31 -0) ➕ `comfy_kitchen/backends/cuda/draft/attention/ada/instantiations/inst_k64_d128_int8_dense.cu` (+16 -0) ➕ `comfy_kitchen/backends/cuda/draft/attention/ada/instantiations/inst_k64_d64.cu` (+31 -0) ➕ `comfy_kitchen/backends/cuda/draft/attention/ada/instantiations/inst_k64_d64_int8_dense.cu` (+16 -0) ➕ `comfy_kitchen/backends/cuda/draft/attention/ada/instantiations/inst_q128_k64_d128.cu` (+19 -0) ➕ `comfy_kitchen/backends/cuda/draft/attention/ada/instantiations/inst_q128_k64_d128_fp16.cu` (+13 -0) ➕ `comfy_kitchen/backends/cuda/draft/attention/ada/instantiations/inst_q128_k64_d128_int8.cu` (+13 -0) ➕ `comfy_kitchen/backends/cuda/draft/attention/ada/instantiations/inst_q128_k64_d64.cu` (+19 -0) ➕ `comfy_kitchen/backends/cuda/draft/attention/ada/instantiations/inst_q128_k64_d64_fp16.cu` (+13 -0) ➕ `comfy_kitchen/backends/cuda/draft/attention/ada/instantiations/inst_q128_k64_d64_int8.cu` (+13 -0) ➕ `comfy_kitchen/backends/cuda/draft/attention/ada/k64_attention_decl.cuh` (+98 -0) _...and 73 more files_ </details> ### 📄 Description ### Summary Adds `comfy_kitchen.draft_attention`, a DraftMap-based block-sparse self-attention operation for packed video sequences, with a full-precision eager reference and native CUDA executors for NVIDIA GPUs from SM89 onward. The implementation is adapted from [`anemoi-project/anemoi`](https://github.com/anemoi-project/anemoi) revision `bcb18202a4afd006749700370634de4819f9cc19`. It consolidates option validation, layout planning, backend dispatch, native availability checks, and the CUDA executor behind one public Comfy Kitchen operation. ### What changed - Adds the public APIs: - `comfy_kitchen.draft_attention(...)` - `comfy_kitchen.draft_attention_is_available(device)` - Registers eager and CUDA implementations with capability-aware dispatch and constraint metadata. - Adds a device-portable full-precision eager reference for CPU and unsupported accelerators/devices. - Adds shared validation and packed video-layout construction without persistent GPU tensor caches. - Adds the migrated native CUDA plan, routing, draft, preparation, attention, and output-assembly stages. - Serves capabilities `[89, 120)` with the Ada kernel set; SM120 uses the SM120a executor. - Adds build-system integration and backend capability reporting. - Adds `tests/test_draft_attention.py` covering the public contract, option and layout validation, dispatch, availability, eager correctness, native stages, and CUDA reference execution. ### Public behavior Inputs are equal-shaped BTHD fp16/bf16 tensors with batch size one plus a packed video shape and prefix length. The default profile is Q64 with 80% sparsity and an INT8 middle phase. The native path is selected only when the complete executor was compiled for the current CUDA device and the call satisfies its constraints. CPU, unsupported accelerators, and CUDA devices without a matching compiled executor fall through to the full-precision eager reference instead of failing dispatch. ### Benchmarks #### End-to-end protocol MiniMax H3 video generation through ComfyUI, with [Yusux/ComfyUI@e65ddf0](https://github.com/Yusux/ComfyUI/commit/e65ddf09324c1a5436288762bd225f07cea8c711) and [Yusux/comfy-kitchen-anemoi@2817e13](https://github.com/Yusux/comfy-kitchen/tree/2817e131b4cdf45710a374d2f1d14d836a73a772) using: - `minimax_h3_fl2va_pruned_fp8_scaled.safetensors` (UNet) - `qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors` (text encoder) - `minimax_h3_video_vae_fp16.safetensors` + `minimax_h3_audio_vae_fp32.safetensors` (VAEs) These measurements were collected before the Anemoi-to-Draft rename. The tables use the current Draft name for the same implementation; E2E benchmarks were not rerun after the naming-only changes. All runs use **1344×768, 124 frames, 24 FPS, and 50 sampling steps**, with the `res_multistep` sampler and `simple` scheduler. The three scenes are city (seed 0), ocean (seed 0), and forest (seed 1234), with identical prompts across systems. Each scene/configuration has one six-step warmup followed by one 50-step formal run on a persistent server; warmup time is excluded. Reported E2E wall time runs from prompt submission to completion and includes sampling, video/audio VAE decode, and video saving. Completion is polled every two seconds. All sparse profiles use Q64 and `start_percent=0.2`. This setting is converted to a sigma threshold rather than an integer step count. In all 27 formal sparse runs, Draft attention starts after the initial 11-step dense phase. Drop denotes the configured sparsity ratio; NVFP4/INT8 percentages describe the precision split of retained sparse work, not the model weights. The attention video grid is `(37, 24, 42)`, with 56 heads, head dimension 128, and 473–479 prefix tokens (depending on the prompt). #### Cross-architecture E2E: Q64, Drop80, 100% INT8 Means across the three scenes (one formal run per scene/configuration). E2E speedups below are Dense mean time divided by sparse mean time: | GPU | Architecture | PyTorch | Dense E2E (s) | Draft E2E (s) | Speedup vs Dense | RGB PSNR vs Dense (dB) | |---|---|---|---:|---:|---:|---:| | RTX 4090 (48 GB) | SM89 | 2.12.1+cu130 | 1037.20 | 615.58 | 1.685× | 20.46 | | H200 | SM90 | 2.12.1+cu130 | 392.58 | 236.43 | 1.660× | 20.51 | | RTX 5090 | SM120 | 2.11.0+cu130 | 704.84 | 388.42 | 1.815× | 20.88 | SM89 and SM90 use the native Ada-family INT8 executor; the mixed NVFP4 attention profiles below were tested only on SM120. PSNR is computed over decoded RGB frames against each scene's Dense output on the **same GPU**, then averaged across scenes. It measures output similarity, not perceptual quality or cross-GPU numerical equivalence; audio is not scored. #### RTX 5090 E2E: seven Q64 sparse profiles All seven profiles use INT8 prefix K/V and prefix queries. | Drop | Retained NVFP4 / INT8 | Mean E2E (s) | Speedup vs Dense | RGB PSNR vs Dense (dB) | |---|---|---:|---:|---:| | — | Dense reference | 704.84 | 1.000× | — | | 80% | 0% / 100% | 388.42 | 1.815× | 20.88 | | 60% | 60% / 40% | 413.25 | 1.706× | 21.22 | | 60% | 75% / 25% | 410.46 | 1.717× | 20.60 | | 70% | 60% / 40% | 399.79 | 1.763× | 20.81 | | 70% | 75% / 25% | 397.17 | 1.775× | 19.83 | | 80% | 60% / 40% | 386.47 | 1.824× | 20.04 | | 80% | 75% / 25% | **383.10** | **1.840×** | 19.44 | At matched Drop80, NV75/I25 reduces mean wall time by 5.32 s versus pure INT8 (about 1.014× additional speedup). Drop60/70 retain more sparse work and remain slower than Drop80 INT8. #### RTX 5090 Q64 preparation / mainloop breakdown Separate operator profiling at the city attention geometry: `B=1, T=37775, H=56, D=128`, video grid `(37, 24, 42)`, prefix length 479. Values are means of three CUDA-profiler captures per configuration. These are **kernel-duration sums for preparation and native attention mainloop only**, not full operator wall time and not the 50-step E2E measurement. Other stages and host overhead are excluded; sums below are rounded from unrounded values. | Drop | Retained NVFP4 / INT8 | Preparation (ms) | Mainloop (ms) | Prep + mainloop (ms) | |---|---|---:|---:|---:| | 80% | 0% / 100% | 3.393 | 15.046 | 18.439 | | 60% | 60% / 40% | 4.479 | 24.583 | 29.062 | | 60% | 75% / 25% | 4.495 | 23.369 | 27.864 | | 70% | 60% / 40% | 4.463 | 18.863 | 23.325 | | 70% | 75% / 25% | 4.418 | 17.698 | 22.115 | | 80% | 60% / 40% | 4.395 | 12.881 | 17.276 | | 80% | 75% / 25% | 4.413 | 12.377 | 16.790 | Mixed NVFP4/INT8 adds approximately 1 ms of preparation, but at matched Drop80 its mainloop saving offsets that cost: prep + mainloop falls by 6.31% for NV60/I40 and 8.94% for NV75/I25. This component-level saving is diluted in E2E by the dense-first sampling portion and work outside sparse attention. ### Provenance and licensing The migrated implementation is adapted from the Apache-2.0 Anemoi source at revision `bcb18202a4afd006749700370634de4819f9cc19`. SPDX headers, Ada/Blackwell license files, Blackwell provenance, and per-file upstream attribution are preserved. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

2026-10-26 22:21:29 +00:00

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#196

No description provided.

Rows
Columns