[PR #88] [MERGED] Add swiglu input_act and partial-rotary rot_dim for rms_rope #134

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

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

Base: main ← Head: kernel_dev

📝 Commits (3)

967c098 Add swiglu input_act and partial-rotary rot_dim for rms_rope
22b3afa Fixes
8ce7286 ruff fix

📊 Changes

13 files changed (+492 additions, -64 deletions)

View changed files

📝 comfy_kitchen/__init__.py (+12 -5)
📝 comfy_kitchen/_rope_utils.py (+37 -3)
📝 comfy_kitchen/backends/_activations.py (+16 -3)
📝 comfy_kitchen/backends/cuda/__init__.py (+34 -14)
📝 comfy_kitchen/backends/cuda/dlpack_bindings.cpp (+20 -7)
➕ comfy_kitchen/backends/cuda/input_act_codes.h (+15 -0)
📝 comfy_kitchen/backends/cuda/ops/int8_linear.cu (+54 -17)
📝 comfy_kitchen/backends/cuda/ops/rms_rope.cu (+40 -7)
📝 comfy_kitchen/backends/eager/rope.py (+20 -8)
📝 comfy_kitchen/backends/hip/__init__.py (+15 -0)
📝 comfy_kitchen/backends/triton/rms_rope.py (+11 -0)
📝 tests/test_int8_input_act.py (+97 -0)
📝 tests/test_rms_rope.py (+121 -0)

📄 Description

Two additions for packed single-stream DiTs, plus build fixes:

int8_linear input_act gains "swiglu": the quantizer reads a [gate | up] row
(2*K wide) and absorbs silu(gate) * up on the way in, so the MLP intermediate
never round-trips through HBM.
rms_rope_split_half(_) gains rot_dim: partial rotary that rotates only a
head-dim prefix (pairs (i, i + rot_dim/2)) while the RMS norm spans the full
head dim. In-place now also accepts q/k slices of a fused QKV projection
(interleaved views were rejected by the conservative overlap check).
And building on Windows surfaced two MSVC 14.42 C3495 errors (constexpr locals captured into nested
dispatch lambdas) that broke the CUDA build, should be behavior-neutral fix.

Tests:

swiglu cases added to test_int8_input_act.py
TestPartialRotary added to test_rms_rope.py

_{🔄 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/88 **Author:** [@kijai](https://github.com/kijai) **Created:** 7/31/2026 **Status:** ✅ Merged **Merged:** 7/31/2026 **Merged by:** [@comfyanonymous](https://github.com/comfyanonymous) **Base:** `main` ← **Head:** `kernel_dev` --- ### 📝 Commits (3) - [`967c098`](https://github.com/Comfy-Org/comfy-kitchen/commit/967c098c5cb9b125ca7d84a8a6e404dc801f8be3) Add swiglu input_act and partial-rotary rot_dim for rms_rope - [`22b3afa`](https://github.com/Comfy-Org/comfy-kitchen/commit/22b3afab7401568b1ff1c6a708a835ebbdbc5864) Fixes - [`8ce7286`](https://github.com/Comfy-Org/comfy-kitchen/commit/8ce7286d5bcf338833b2d42bde0272660dfb4e52) ruff fix ### 📊 Changes **13 files changed** (+492 additions, -64 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/__init__.py` (+12 -5) 📝 `comfy_kitchen/_rope_utils.py` (+37 -3) 📝 `comfy_kitchen/backends/_activations.py` (+16 -3) 📝 `comfy_kitchen/backends/cuda/__init__.py` (+34 -14) 📝 `comfy_kitchen/backends/cuda/dlpack_bindings.cpp` (+20 -7) ➕ `comfy_kitchen/backends/cuda/input_act_codes.h` (+15 -0) 📝 `comfy_kitchen/backends/cuda/ops/int8_linear.cu` (+54 -17) 📝 `comfy_kitchen/backends/cuda/ops/rms_rope.cu` (+40 -7) 📝 `comfy_kitchen/backends/eager/rope.py` (+20 -8) 📝 `comfy_kitchen/backends/hip/__init__.py` (+15 -0) 📝 `comfy_kitchen/backends/triton/rms_rope.py` (+11 -0) 📝 `tests/test_int8_input_act.py` (+97 -0) 📝 `tests/test_rms_rope.py` (+121 -0) </details> ### 📄 Description Two additions for packed single-stream DiTs, plus build fixes: - int8_linear input_act gains "swiglu": the quantizer reads a [gate | up] row (2*K wide) and absorbs silu(gate) * up on the way in, so the MLP intermediate never round-trips through HBM. - rms_rope_split_half(_) gains rot_dim: partial rotary that rotates only a head-dim prefix (pairs (i, i + rot_dim/2)) while the RMS norm spans the full head dim. In-place now also accepts q/k slices of a fused QKV projection (interleaved views were rejected by the conservative overlap check). - And building on Windows surfaced two MSVC 14.42 C3495 errors (constexpr locals captured into nested dispatch lambdas) that broke the CUDA build, should be behavior-neutral fix. Tests: - swiglu cases added to test_int8_input_act.py - TestPartialRotary added to test_rms_rope.py --- <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:42 +00:00