[IvenHsu01]

Gated residual for reference-image tokens (index_timestep_zero) in image-edit DiTs: ck.gated_residual_segmented(residual, x, gate, split) returns residual + x * gate[i] for x and residual [B, L, D], with i = 0 for tokens [0, split) and i = 1 for [split, L); ck.gated_residual_segmented_ writes it into residual in place. gate is [2, B, D] (or [2, 1, D]), a view of each model's modulation output: Flux gate.transpose(0, 1) of [B, 2, D], Qwen-Image gate.unflatten(0, (2, -1)) of [2B, D], Krea2 gate.squeeze(1).unflatten(0, (2, -1)) of [2B, 1, D]. Used by Qwen-Image _apply_gate, Krea2 SingleStreamBlock.forward and the Flux / Flux.2 x += apply_mod(..., gate, None, modulation_dims) lines in DoubleStreamBlock.forward / SingleStreamBlock.forward, next to adaln_segmented (#307); ComfyUI PR to follow.

ComfyUI currently runs one multiply per token segment (Qwen-Image then joins them with torch.cat(..., dim=1)) and the residual add as separate kernels, in every block of every step. The Triton kernel holds one token row in registers per program, reads the gate rows through their batch strides (no copy of the views) and stores residual + x * gate[i] computed in fp32 and rounded once, into a new tensor or in place. fp16 / bf16 / fp32, D up to 8192; inputs the kernel does not cover run eager. The launch is fixed, not autotuned, because autotuning would re-run an in-place launch.

Speedup over the current path per call through the public op, bf16, at the Krea2 and Qwen-Image shapes captured from real runs and the Flux.2 Klein 9B shapes derived from its config; every call was served by triton, and the max and mean error against fp32 are no larger than the current path's:

gate, bf16	shapes	RTX 5090	Radeon 8060S (RDNA3.5)
Krea2	1	1.72-1.73x	1.58-1.60x
Qwen-Image	1	2.30x	2.27-2.28x
Flux.2, in place, step 1	2	1.70-1.72x	1.57-1.59x
Flux.2, in place, KV cache	2	1.01-1.12x	1.44-1.46x

Per shape (times from the first of two pairs, speedups from both pairs)

Krea2 (1, 8414, 6144) split 4270, Qwen-Image (1, 12336, 3072) split 4096, Flux.2 Klein 9B double / single blocks (1, 12168, 4096) split 4056 and (1, 12680, 4096) split 4568 on the first step, (1, 4056, 4096) split 4056 / (1, 4568, 4096) split 4568 after the KV cache step. Krea2 and Qwen-Image use gated_residual_segmented, Flux.2 uses gated_residual_segmented_ (in place).

per call, bf16 (current path -> triton)	RTX 5090	Radeon 8060S
Krea2 gate	340.5 -> 197.2 us (1.72-1.73x)	2358.2 -> 1474.4 us (1.58-1.60x)
Qwen-Image gate	325.4 -> 141.2 us (2.30x)	2480.1 -> 1090.8 us (2.27-2.28x)
Flux.2 double gate, step 1	324.5 -> 189.7 us (1.70-1.71x)	2241.8 -> 1411.7 us (1.59x)
Flux.2 single gate, step 1	340.4 -> 198.0 us (1.72x)	2322.8 -> 1472.8 us (1.57-1.58x)
Flux.2 double gate, KV cache	37.0 -> 36.8 us (1.01-1.04x)	785.4 -> 537.7 us (1.45-1.46x)
Flux.2 single gate, KV cache	41.1 -> 36.8 us (1.12x)	872.5 -> 602.9 us (1.44-1.45x)

Windows: RTX 5090 with torch 2.14.1+cu130, triton-windows 3.9.0 and the CUDA extension, Radeon 8060S with ROCm 10.1 torch 2.15.0a0, triton-windows 3.8.0 and the HIP extension, both from the 0.2.41 wheel. Each arm in a fresh process, ABAB, median of 5 x 100 back-to-back calls after 20 warmup calls. Measured on 487f577; the head has the same code (later changes only edit comments and docstrings, and the history was rebuilt so each commit holds only this op). tests/test_gated_residual_segmented.py (32 tests), test_backends.py and test_constraints.py pass on both GPUs (2 CUDA-only skips on the Radeon 8060S).

End to end in ComfyUI on the Radeon 8060S (ComfyUI fb2b26f, --enable-triton-backend in both arms, COMFYUI_ENABLE_MIOPEN=1, a fresh server per arm, ABAB): the Qwen-Image Edit 2511 int8 template at its own settings (1392x752, 40 steps), with adaln_segmented (#307) and gated_residual_segmented (#310) swapped in together on the segmented image path of QwenImageTransformerBlock (9600 + 9600 calls per prompt, all served by triton), takes 682.99 -> 665.45 s and 683.31 -> 664.64 s per prompt, 2.57% and 2.73% less time. The output is the same picture as the stock run, at 32.78 dB PSNR, since the fused chain rounds once in fp32. Measured with the code of 53fe266 and 487f577.

Not covered: end to end on the RTX 5090 and for Krea2 and Flux.2, the tests on Linux.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts