[diffusion] fix: keep residual_gate_add on the JIT CUDA path for cont…

[niehen6174]

Motivation

#41305 replaced the shared residual_gate_add JIT CUDA operator with Triton kernels, benchmarked on SANA-Video on H200 where it gained 1.6–1.7% E2E. That PR's profiling confirmed the branch executing there is _rga_transposed.

The kernel is also shared by Ideogram, Flux-2, Qwen-Image 2.1, ERNIE-Image and LTX-2, which call it on contiguous layouts. There the Triton path costs 13–15.6 µs more host time per call. The cost is flat in tensor size, so it is per-call dispatch overhead, not slower math — negligible when bandwidth-bound, dominant when launch-bound.

LTX-2 is launch-bound: ~800 calls per denoise step at (1, 384, 4096) and (1, 26, 2048), from 8 call sites across 48 layers. That adds ~8 ms per step and moved ltx_2_two_stage_t2v and ltx_2_3_two_stage_ti2v_2gpus from ~1.12× to 1.20–1.23× of their h100.json baselines. The pr_test limit is 1.20×, so about half of CI runs now cross it.

Found by bisecting 315 commits on 2× H200. Reverting this one file on current main removes the regression entirely.

Modifications

One file, +9/−4.

1. Take the Triton path only for the transposed-dense layout it was benchmarked on; contiguous layouts go back to the JIT CUDA kernel. The SANA-Video path is untouched.
2. Drop two redundant can_use_residual_gate_add_cuda calls — it ran 3× per invocation at ~2.7 µs each (pre-[KDA+Kimi K3] Speed up SANA-Video residual gate add on H200 #41305 ran it twice).
3. Call the custom op directly from residual_gate_add, which has already run the same predicate. residual_gate_add_cuda stays public and keeps validating.

Accuracy Tests

The CUDA and Triton backends were compared on 37 configurations (bf16/fp16, six shapes, full / broadcast / per-token gates, plus transposed-dense) and are bit-identical on all of them (torch.equal). Both already claimed bit-exact rounding, so routing between them by layout cannot change output.

Speed Tests and Profiling

2× H200, dependencies pinned, sglang commit the only variable.

Host time per call, bf16, at the shapes LTX-2 actually issues (confirmed by counting dispatch branches over a real run):

shape	gate	#41305	this PR	pre-#41305
(1, 26, 2048)	broadcast	36.50 µs	23.80 µs	23.46 µs
(1, 384, 4096)	broadcast	39.18 µs	23.92 µs	23.43 µs
(1, 1536, 4096)	full	34.43 µs	19.91 µs	19.19 µs

ltx_2_two_stage_t2v avg denoise step, three runs each:

build	runs	median
544e46c080 (current main)	164.7 / 162.5 / 161.6	162.5 ms
main + this PR	158.2 / 158.1 / 157.4	158.1 ms
main + whole file reverted	157.1 / 156.8 / 157.0	157.0 ms
fc9e1c8d29 (before #41305)	156.3 / 157.9 / 154.7	156.3 ms

Recovers 4.4 ms of the 6.2 ms regression. A ~1.1 ms gap to the full revert remains: 672 of 24000 calls (gate same shape as residual) still fall through to the eager path. Can follow up separately.

The microbenchmark covers contiguous layouts only — I did not measure whether Triton also wins there at SANA-Video's sizes. If it does, the dispatch can be widened on the back of a benchmark covering the other models sharing this kernel; test/registered/kernels/benchmark/diffusion/bench_residual_gate_add.py already has workloads for most of them.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests. — no new behaviour; existing kernel tests plus the bit-exactness check above apply
• Update documentation according to Write documentations. — no user-facing change
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

cc @BBuf

---

CI States

Latest PR Test (Base): ⏳ Run #37088752701
Latest PR Test (Extra): ❌ Run #37088767953
Latest PR Test (AMD ROCm 10): ❌ Run #37088752682