[PR #167] [MERGED] MiniMax H3 VAE kernels: fused encoder pad/norm, fp16-accumulate conv3d and GEMM, int8 residual epilogue #195

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

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

Base: main ← Head: minimax_vae

📝 Commits (10+)

ad9a4a8 initial int8 sol-attn implemenation
75113d3 Merge remote-tracking branch 'upstream/main' into sol_attn
6716dfa Merge remote-tracking branch 'upstream/main' into sol_attn
c27d802 sol_attn: centroid-tail routing (1.38x), u8 P quantization, and key_bias masks
79a7532 Fixes
d7fc603 Merge remote-tracking branch 'upstream/main' into sol_attn
883e710 ruff
c04ef20 Add reuse_qkv_memory to reduce peak VRAM
85c4142 Merge remote-tracking branch 'upstream/main' into sol_attn
b17db5d Optimize routing

📊 Changes

27 files changed (+2914 additions, -558 deletions)

View changed files

📝 comfy_kitchen/__init__.py (+97 -5)
📝 comfy_kitchen/backends/_activations.py (+29 -1)
📝 comfy_kitchen/backends/cuda/CMakeLists.txt (+3 -0)
📝 comfy_kitchen/backends/cuda/__init__.py (+333 -42)
📝 comfy_kitchen/backends/cuda/dlpack_bindings.cpp (+321 -27)
📝 comfy_kitchen/backends/cuda/input_act_codes.h (+2 -2)
📝 comfy_kitchen/backends/cuda/ops/convrot_w4a4.cu (+14 -11)
➕ comfy_kitchen/backends/cuda/ops/cutlass_conv3d_fp16.cu (+146 -0)
➕ comfy_kitchen/backends/cuda/ops/cutlass_gemm_common.cuh (+259 -0)
➕ comfy_kitchen/backends/cuda/ops/cutlass_gemm_fp16.cu (+246 -0)
📝 comfy_kitchen/backends/cuda/ops/cutlass_gemm_int8.cu (+174 -433)
➕ comfy_kitchen/backends/cuda/ops/group_norm_pad3d.cu (+226 -0)
📝 comfy_kitchen/backends/cuda/ops/int8_linear.cu (+81 -17)
📝 comfy_kitchen/backends/cuda/ops/w4a8_gemm.cu (+5 -2)
📝 comfy_kitchen/backends/eager/__init__.py (+33 -0)
➕ comfy_kitchen/backends/eager/conv3d.py (+44 -0)
➕ comfy_kitchen/backends/eager/group_norm_pad3d.py (+67 -0)
📝 comfy_kitchen/backends/eager/quantization.py (+38 -3)
📝 comfy_kitchen/backends/hip/__init__.py (+15 -2)
📝 comfy_kitchen/backends/triton/quantization.py (+13 -2)

...and 7 more files

📄 Description

Kernels to speedup the MiniMax-H3 video VAE. Requires core side PR to use:

https://github.com/Comfy-Org/ComfyUI/pull/16187

Encoder (conv)

group_norm_silu_pad3d: per-frame GroupNorm + SiLU + reflect/causal padding in one
pass, NDHWC output so cuDNN runs without layout conversions. Lossless, always on.
fp16_conv3d: CUTLASS fp16-accumulate conv3d with bias and residual in the epilogue.
Only used under --fast fp16_accumulation

Decoder (ViT)

fp16_linear: CUTLASS fp16-accumulate GEMM with a fused residual (addcmul) epilogue,
launches too small to fill the GPU are handed back to cuBLAS.
int8_linear: fused residual epilogue, rms_norm/swiglu folded into the activation
quantizer, wave-quantization guard between the 128x256/128x128 tiles, quantizer
block sizes tuned per K. Bias is read in the output dtype across the GEMM family.

Measured on a 5090 at 1344x768, 362 frames (end-to-end comfy.sd.VAE calls):
encode fp16 46.6 -> 32.0 s, encode with the flag 47.2 -> 21.2 s,
decode fp16 34.6 -> 31.1 s, decode with the flag 27.7 -> 23.2 s, decode int8 18.7 -> 12.9 s.

Quality: int8 decode 67.7 dB and fp16-accum encode 68 dB against the exact path
(the VAE's own reconstruction floor is ~38 dB).

_{🔄 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/167 **Author:** [@kijai](https://github.com/kijai) **Created:** 9/8/2026 **Status:** ✅ Merged **Merged:** 9/13/2026 **Merged by:** [@comfyanonymous](https://github.com/comfyanonymous) **Base:** `main` ← **Head:** `minimax_vae` --- ### 📝 Commits (10+) - [`ad9a4a8`](https://github.com/Comfy-Org/comfy-kitchen/commit/ad9a4a80029fe4e0bc787247b6beaaf29c4a564e) initial int8 sol-attn implemenation - [`75113d3`](https://github.com/Comfy-Org/comfy-kitchen/commit/75113d36fdf909280bab89355bb4cae508d45e3b) Merge remote-tracking branch 'upstream/main' into sol_attn - [`6716dfa`](https://github.com/Comfy-Org/comfy-kitchen/commit/6716dfa59cd5227e434c748c799a58cdf835b328) Merge remote-tracking branch 'upstream/main' into sol_attn - [`c27d802`](https://github.com/Comfy-Org/comfy-kitchen/commit/c27d802f1448e2bbce2721cab23ef0549b758ba2) sol_attn: centroid-tail routing (1.38x), u8 P quantization, and key_bias masks - [`79a7532`](https://github.com/Comfy-Org/comfy-kitchen/commit/79a75320aa71235dae671fee6904a18c24b08e90) Fixes - [`d7fc603`](https://github.com/Comfy-Org/comfy-kitchen/commit/d7fc60382a5381c3f174e27dbe810b1b05c38da7) Merge remote-tracking branch 'upstream/main' into sol_attn - [`883e710`](https://github.com/Comfy-Org/comfy-kitchen/commit/883e710f8a779c7766171ed6e2b07c1e8e032836) ruff - [`c04ef20`](https://github.com/Comfy-Org/comfy-kitchen/commit/c04ef206f3b84cedb7da26414547daa0fe789bea) Add reuse_qkv_memory to reduce peak VRAM - [`85c4142`](https://github.com/Comfy-Org/comfy-kitchen/commit/85c41421f710a2e4eb415ee2eb457ba11bf68956) Merge remote-tracking branch 'upstream/main' into sol_attn - [`b17db5d`](https://github.com/Comfy-Org/comfy-kitchen/commit/b17db5d7041d5c6fdccb501952fceaeb189c10a2) Optimize routing ### 📊 Changes **27 files changed** (+2914 additions, -558 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/__init__.py` (+97 -5) 📝 `comfy_kitchen/backends/_activations.py` (+29 -1) 📝 `comfy_kitchen/backends/cuda/CMakeLists.txt` (+3 -0) 📝 `comfy_kitchen/backends/cuda/__init__.py` (+333 -42) 📝 `comfy_kitchen/backends/cuda/dlpack_bindings.cpp` (+321 -27) 📝 `comfy_kitchen/backends/cuda/input_act_codes.h` (+2 -2) 📝 `comfy_kitchen/backends/cuda/ops/convrot_w4a4.cu` (+14 -11) ➕ `comfy_kitchen/backends/cuda/ops/cutlass_conv3d_fp16.cu` (+146 -0) ➕ `comfy_kitchen/backends/cuda/ops/cutlass_gemm_common.cuh` (+259 -0) ➕ `comfy_kitchen/backends/cuda/ops/cutlass_gemm_fp16.cu` (+246 -0) 📝 `comfy_kitchen/backends/cuda/ops/cutlass_gemm_int8.cu` (+174 -433) ➕ `comfy_kitchen/backends/cuda/ops/group_norm_pad3d.cu` (+226 -0) 📝 `comfy_kitchen/backends/cuda/ops/int8_linear.cu` (+81 -17) 📝 `comfy_kitchen/backends/cuda/ops/w4a8_gemm.cu` (+5 -2) 📝 `comfy_kitchen/backends/eager/__init__.py` (+33 -0) ➕ `comfy_kitchen/backends/eager/conv3d.py` (+44 -0) ➕ `comfy_kitchen/backends/eager/group_norm_pad3d.py` (+67 -0) 📝 `comfy_kitchen/backends/eager/quantization.py` (+38 -3) 📝 `comfy_kitchen/backends/hip/__init__.py` (+15 -2) 📝 `comfy_kitchen/backends/triton/quantization.py` (+13 -2) _...and 7 more files_ </details> ### 📄 Description Kernels to speedup the MiniMax-H3 video VAE. Requires core side PR to use: https://github.com/Comfy-Org/ComfyUI/pull/16187 Encoder (conv) - group_norm_silu_pad3d: per-frame GroupNorm + SiLU + reflect/causal padding in one pass, NDHWC output so cuDNN runs without layout conversions. Lossless, always on. - fp16_conv3d: CUTLASS fp16-accumulate conv3d with bias and residual in the epilogue. Only used under --fast fp16_accumulation Decoder (ViT) - fp16_linear: CUTLASS fp16-accumulate GEMM with a fused residual (addcmul) epilogue, launches too small to fill the GPU are handed back to cuBLAS. - int8_linear: fused residual epilogue, rms_norm/swiglu folded into the activation quantizer, wave-quantization guard between the 128x256/128x128 tiles, quantizer block sizes tuned per K. Bias is read in the output dtype across the GEMM family. Measured on a 5090 at 1344x768, 362 frames (end-to-end comfy.sd.VAE calls): encode fp16 46.6 -> 32.0 s, encode with the flag 47.2 -> 21.2 s, decode fp16 34.6 -> 31.1 s, decode with the flag 27.7 -> 23.2 s, decode int8 18.7 -> 12.9 s. Quality: int8 decode 67.7 dB and fp16-accum encode 68 dB against the exact path (the VAE's own reconstruction floor is ~38 dB). --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

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