[PR #86] [MERGED] rms_adaln and input_act kernels #135

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

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

Base: main ← Head: ltx2_kernels

📝 Commits (3)

b7f20c6 rms_adaln and input_act kernels
5ea7a08 For correctnes
24e680e Merge branch 'main' into ltx2_kernels

📊 Changes

16 files changed (+759 additions, -125 deletions)

View changed files

📝 comfy_kitchen/__init__.py (+31 -0)
➕ comfy_kitchen/backends/_activations.py (+45 -0)
📝 comfy_kitchen/backends/cuda/__init__.py (+64 -12)
📝 comfy_kitchen/backends/cuda/dlpack_bindings.cpp (+48 -4)
📝 comfy_kitchen/backends/cuda/dtype_dispatch.cuh (+14 -0)
📝 comfy_kitchen/backends/cuda/ops/adaln.cu (+59 -47)
📝 comfy_kitchen/backends/cuda/ops/int8_linear.cu (+56 -51)
📝 comfy_kitchen/backends/cuda/utils.cuh (+11 -0)
📝 comfy_kitchen/backends/eager/__init__.py (+10 -1)
📝 comfy_kitchen/backends/eager/adaln.py (+24 -0)
📝 comfy_kitchen/backends/eager/quantization.py (+7 -1)
📝 comfy_kitchen/backends/triton/__init__.py (+11 -1)
📝 comfy_kitchen/backends/triton/adaln.py (+27 -8)
📝 comfy_kitchen/backends/triton/quantization.py (+3 -0)
📝 tests/test_adaln.py (+173 -0)
➕ tests/test_int8_input_act.py (+176 -0)

📄 Description

Add rms_adaln and int8_linear(input_act=...)

Two fusions for RMSNorm-style DiTs (LTX), reusing existing kernels rather
than adding new ones. Tested in practice that when applied, this gives small but free speed boost for LTX2 inference.

rms_adaln(x, scale, shift, eps) — rmsnorm(x) * (1 + scale) + shift. The
existing adaln is LayerNorm-based and not substitutable. RMSNorm is LayerNorm
with the mean pinned to zero, so adaln.cu is templated on kSubtractMean and
both share one kernel. All three backends. 3.0x over eager at 14080x4096.

int8_linear(..., input_act="gelu_tanh") — folds an MLP's activation into
the ConvRot INT8 quantizer, so linear(act(proj(x))) stops round-tripping act's
output through HBM. Costs +0.02 ms on the quantizer, removes a 0.25 ms kernel,
and drops FF peak VRAM 36%. One optional parameter, not a new op; paths that
can't fuse apply the activation eagerly, so all routes agree.

7.1% end to end on LTX2 in ComfyUI (1.84 -> 1.97 it/s, int8+ConvRot).

Notes

modulation_row moved to utils.cuh; must match _modulation.py, where a
divide-vs-modulo mix-up is silent and only diverges for batch > 1.
Activation helper and code mapping in backends/_activations.py, mirroring
_modulation.py, keeping the enum in sync with int8_linear.cu.
Both fused paths are slightly more accurate than the chains they replace —
neither rounds the intermediate to the input dtype.

Testing

121 tests (62 new in test_adaln.py, 22 in test_int8_input_act.py): broadcast
at batch 1/2/3, per-token modulation, odd and large D, both CUDA kernel paths,
non-contiguous inputs, and every fallback route. adaln output verified
byte-identical across 3 dtypes x 10 shapes.

_{🔄 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/86 **Author:** [@kijai](https://github.com/kijai) **Created:** 7/27/2026 **Status:** ✅ Merged **Merged:** 7/27/2026 **Merged by:** [@comfyanonymous](https://github.com/comfyanonymous) **Base:** `main` ← **Head:** `ltx2_kernels` --- ### 📝 Commits (3) - [`b7f20c6`](https://github.com/Comfy-Org/comfy-kitchen/commit/b7f20c64d5f16427384ff6a4859f5cd332e7737f) rms_adaln and input_act kernels - [`5ea7a08`](https://github.com/Comfy-Org/comfy-kitchen/commit/5ea7a08bc7c1369c7e5bd9f5f970c412e7c0f066) For correctnes - [`24e680e`](https://github.com/Comfy-Org/comfy-kitchen/commit/24e680e5a0e216439ea8079648cf53e8959b5348) Merge branch 'main' into ltx2_kernels ### 📊 Changes **16 files changed** (+759 additions, -125 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/__init__.py` (+31 -0) ➕ `comfy_kitchen/backends/_activations.py` (+45 -0) 📝 `comfy_kitchen/backends/cuda/__init__.py` (+64 -12) 📝 `comfy_kitchen/backends/cuda/dlpack_bindings.cpp` (+48 -4) 📝 `comfy_kitchen/backends/cuda/dtype_dispatch.cuh` (+14 -0) 📝 `comfy_kitchen/backends/cuda/ops/adaln.cu` (+59 -47) 📝 `comfy_kitchen/backends/cuda/ops/int8_linear.cu` (+56 -51) 📝 `comfy_kitchen/backends/cuda/utils.cuh` (+11 -0) 📝 `comfy_kitchen/backends/eager/__init__.py` (+10 -1) 📝 `comfy_kitchen/backends/eager/adaln.py` (+24 -0) 📝 `comfy_kitchen/backends/eager/quantization.py` (+7 -1) 📝 `comfy_kitchen/backends/triton/__init__.py` (+11 -1) 📝 `comfy_kitchen/backends/triton/adaln.py` (+27 -8) 📝 `comfy_kitchen/backends/triton/quantization.py` (+3 -0) 📝 `tests/test_adaln.py` (+173 -0) ➕ `tests/test_int8_input_act.py` (+176 -0) </details> ### 📄 Description ## Add rms_adaln and int8_linear(input_act=...) Two fusions for RMSNorm-style DiTs (LTX), reusing existing kernels rather than adding new ones. Tested in practice that when applied, this gives small but free speed boost for LTX2 inference. **`rms_adaln(x, scale, shift, eps)`** — `rmsnorm(x) * (1 + scale) + shift`. The existing `adaln` is LayerNorm-based and not substitutable. RMSNorm is LayerNorm with the mean pinned to zero, so `adaln.cu` is templated on `kSubtractMean` and both share one kernel. All three backends. **3.0x over eager** at 14080x4096. **`int8_linear(..., input_act="gelu_tanh")`** — folds an MLP's activation into the ConvRot INT8 quantizer, so `linear(act(proj(x)))` stops round-tripping act's output through HBM. Costs +0.02 ms on the quantizer, removes a 0.25 ms kernel, and drops FF peak VRAM 36%. One optional parameter, not a new op; paths that can't fuse apply the activation eagerly, so all routes agree. **7.1% end to end** on LTX2 in ComfyUI (1.84 -> 1.97 it/s, int8+ConvRot). ### Notes - `modulation_row` moved to `utils.cuh`; must match `_modulation.py`, where a divide-vs-modulo mix-up is silent and only diverges for batch > 1. - Activation helper and code mapping in `backends/_activations.py`, mirroring `_modulation.py`, keeping the enum in sync with `int8_linear.cu`. - Both fused paths are slightly *more* accurate than the chains they replace — neither rounds the intermediate to the input dtype. ### Testing 121 tests (62 new in `test_adaln.py`, 22 in `test_int8_input_act.py`): broadcast at batch 1/2/3, per-token modulation, odd and large `D`, both CUDA kernel paths, non-contiguous inputs, and every fallback route. `adaln` output verified byte-identical across 3 dtypes x 10 shapes. --- <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:43 +00:00