[IvenHsu01]

Adds ck.adaln_segmented(x, scale, shift, split, eps=1e-6, rms_weight=None, rms_weight_offset=0.0) for the reference-token (index_timestep_zero) modulation of image-edit DiTs, where the tokens before a split index are modulated with one set of scale / shift rows and the reference-image tokens after it with another. It computes norm(x) * (1 + scale[i]) + shift[i] for x [B, L, D], with i = 0 for tokens [0, split) and i = 1 for [split, L). The norm is LayerNorm without affine, or, with rms_weight given, an fp32 RMSNorm with weight rms_weight + rms_weight_offset rounded to x's dtype after the norm (Krea2 stores its RMSNorm scale zero-centered, offset 1.0).

scale / shift are [2, B, D] (or [2, 1, D]), any strides on the first two dims, which each model gets from its modulation output as a view: Flux scale.transpose(0, 1) of [B, 2, D], Qwen-Image and Krea2 scale.unflatten(0, (2, -1)) of [2B, D].

In ComfyUI this runs as a norm followed by in-place updates of per-segment slices, each a separate kernel reading and writing the full activation, before the attention and the MLP of every block of every step:

• Qwen-Image _modulate (comfy/ldm/qwen_image/model.py): F.layer_norm, then torch.addcmul per segment and torch.cat.
• Krea2 SingleStreamBlock.forward (comfy/ldm/krea2/model.py): fp32 RMSNorm, .to(dtype), then mul_(1 + scale).add_(shift) per segment.
• Flux / Flux.2 modulated_norm with modulation_dims (comfy/ldm/flux/layers.py): with modulation_dims set, modulated_norm skips ck.adaln and runs layer_norm, then *= / += per segment.

The segmented gated residual after the attention / MLP of the same blocks (gated_residual_segmented) is in a separate PR: #310.

• eager: the reference sequence, in the Flux / Krea2 op order (the product t * (1 + scale) rounded to x's dtype before the shift is added); it equals those two paths exactly. Qwen-Image's addcmul rounds once, so it can differ from eager by one rounding step.
• triton: one program per token row holds the whole row in registers (block_d = next_power_of_2(D), num_warps 8 from block_d 4096 up, else 4), so x is read once; the norm statistics, the normalize, the optional RMSNorm weight and the scale / shift row of the token's segment are applied in fp32 and rounded once on the store. Segment rows are read through their own batch strides, so the [2, B, D] views are not copied. fp16 / bf16 / fp32 x with modulation of the same dtype and D up to 8192; x needs unit stride on D and is made contiguous otherwise, and empty inputs run eager. This launch was the fastest or within 3% of the fastest of the configs tried (a two-pass loop over 1024 / 2048-column chunks, num_warps autotuning) at every captured shape on the RTX 5090 and the Radeon 8060S.

Speedup over ComfyUI's current path, bf16, per call through the public op (so the op's host-side dispatch is included), at the shapes captured from real runs: 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) / (1, 4568, 4096) after the KV cache step. Times are from the first of two pairs, the speedup is given for both pairs. Every call was served by triton, and at every shape the max and mean error against an fp32 reference are no larger than the current path's:

per call, bf16 (current path -> triton)	RTX 5090	Radeon 8060S
Krea2 (RMSNorm)	965.0 -> 134.5 us (7.17x)	6685.4 -> 1064.6 us (6.27-6.28x)
Qwen-Image	296.8 -> 92.8 us (3.20x)	2151.9 -> 813.2 us (2.65-2.69x)
Flux.2 double, step 1	332.9 -> 129.9 us (2.55-2.56x)	2584.9 -> 995.2 us (2.59-2.60x)
Flux.2 single, step 1	346.6 -> 135.5 us (2.56x)	2853.7 -> 1036.3 us (2.75x)
Flux.2 double, KV cache	64.0 -> 46.4 us (1.38-1.43x)	765.3 -> 418.9 us (1.83-1.85x)
Flux.2 single, KV cache	73.2 -> 45.9 us (1.59-1.62x)	1070.4 -> 452.7 us (2.36x)

Environment: RTX 5090 on Windows (torch 2.14.1+cu130, triton-windows 3.9.0, Python 3.12.16); Radeon 8060S (RDNA3.5) on Windows (ROCm torch 2.15.0a0+rocm10.1.0a20260907, HIP 7.16, triton-windows 3.8.0, Python 3.12.10).

Tests and benchmarks on 53fe266 on both GPUs (the head 076094d after it only edits comments), with prebuilt extensions copied into the source tree: test_adaln_segmented.py 33, test_backends.py 12, test_constraints.py 42 passed on the RTX 5090 (CUDA extension from the 0.2.41 wheel); the same on the Radeon 8060S (HIP extension only), except test_constraints.py 40 passed / 2 skipped (CUDA-only).

Not covered yet: end to end in ComfyUI, on either GPU. ComfyUI integration (Qwen-Image, Krea2, Flux / Flux.2 call sites) is a separate PR; ComfyUI keeps comfy-kitchen's triton backend off unless the server is started with --enable-triton-backend.

[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

Tip

🎯 Customize your risk assessments

Guide how CodeRabbit assesses merge risk and security architecture risk. Highlight sensitive changes and explain your security boundaries using Merge Risk Instructions and Security Risk Instructions in your review settings. Security architecture risk assessments require the Advanced plan or above.

✅ Automatic approvals, with your rules

Let CodeRabbit approve pull requests once required threads are resolved, pre-merge checks pass, and your risk limits are met. Enable Automatic approvals in review settings, with overrides for teams or file paths.