From 3e7628ceab241ca61db084b7398ecf75b8828236 Mon Sep 17 00:00:00 2001 From: niehen6174 <1639206518@qq.com> Date: Fri, 2 Oct 2026 15:43:29 +0000 Subject: [PATCH] [diffusion] fix: keep residual_gate_add on the JIT CUDA path for contiguous layouts #41305 replaced the shared residual_gate_add JIT CUDA operator with Triton kernels and a dispatch heuristic tuned for SANA-Video on H200. The kernel is shared by Ideogram, Flux-2, Qwen-Image 2.1, ERNIE-Image and LTX-2, and for the contiguous layouts those models use the Triton path costs 13-15.6us more host time per call. The extra cost is flat in tensor size, so it is per-call dispatch overhead rather than slower arithmetic, and it dominates whenever the call is launch-bound. LTX-2 is launch-bound here: a denoise step issues roughly 800 calls, at (1, 384, 4096) on the video path and (1, 26, 2048) on the audio path, from eight call sites across 48 layers. That turned into ~8ms per denoise step and pushed ltx_2_two_stage_t2v and ltx_2_3_two_stage_ti2v_2gpus onto the 20% pr_test limit, where about half of all CI runs now cross it. Take the Triton path only for the transposed-dense layout it was benchmarked on and keep the JIT CUDA kernel elsewhere. Also drop two redundant can_use_residual_gate_add_cuda calls: it ran three times per invocation (entry, residual_gate_add_cuda, _residual_gate_add_triton) at ~2.7us each, where the pre-#41305 code ran it twice. Measured on 2x H200, dependencies pinned, sglang commit the only variable. Host time per call, bfloat16, at the shapes LTX-2 actually issues: shape #41305 this patch pre-#41305 (1, 26, 2048) bcast 36.50us 23.80us 23.46us (1, 384, 4096) bcast 39.18us 23.92us 23.43us (1, 1536, 4096) full 34.43us 19.91us 19.19us (1, 384, 4096) full 34.66us 20.35us 20.04us ltx_2_two_stage_t2v avg denoise, three runs each: main 544e46c080 164.7 / 162.5 / 161.6 median 162.5 main + this patch 158.2 / 158.1 / 157.4 median 158.1 main + whole file reverted 157.1 / 156.8 / 157.0 median 157.0 fc9e1c8d29 (before #41305) 156.3 / 157.9 / 154.7 median 156.3 Output is unchanged: the CUDA and Triton backends were compared on 37 configurations (bfloat16 and float16, six shapes, full/broadcast/per-token gates, plus the transposed-dense layout) and are bit-identical on all of them. The SANA-Video shapes keep the Triton path. If it also wins on contiguous layouts at the sizes SANA-Video issues, the dispatch can be widened again on the back of a benchmark that covers the other models sharing this kernel. --- .../kernels/kda_kernels/residual_gate_add_jit.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/python/sglang/kernels/kda_kernels/residual_gate_add_jit.py b/python/sglang/kernels/kda_kernels/residual_gate_add_jit.py index e5892561f8ea..c11b113e7a85 100644 --- a/python/sglang/kernels/kda_kernels/residual_gate_add_jit.py +++ b/python/sglang/kernels/kda_kernels/residual_gate_add_jit.py @@ -269,8 +269,8 @@ def _rga_transposed( def _residual_gate_add_triton(residual, update, gate): - if not can_use_residual_gate_add_cuda(residual, update, gate): - raise RuntimeError("unsupported input for residual_gate_add CUDA") + # Both callers validate the inputs first; repeating the predicate here cost + # ~2.7us of host time on every call. out = torch.empty_strided( residual.shape, residual.stride(), dtype=residual.dtype, device=residual.device ) @@ -420,7 +420,12 @@ def _residual_gate_add_custom_op( residual: torch.Tensor, update: torch.Tensor, gate: torch.Tensor ) -> torch.Tensor: with torch.cuda.device(residual.device): - return _residual_gate_add_triton(residual, update, gate) + # Take the Triton path only for the transposed-dense layout it was + # benchmarked on. For the contiguous layouts every other diffusion model + # uses, the existing JIT CUDA kernel is faster at every measured shape. + if _is_transposed_dense_residual(residual, update, gate): + return _residual_gate_add_triton(residual, update, gate) + return _residual_gate_add_cuda_impl(residual, update, gate) def _gate_mode(residual: torch.Tensor, gate: torch.Tensor) -> int: @@ -513,7 +518,7 @@ def residual_gate_add( and can_use_residual_gate_add_cuda(residual, update, gate) ): try: - return residual_gate_add_cuda(residual, update, gate) + return _residual_gate_add_custom_op(residual, update, gate) except Exception as exc: if torch.compiler.is_compiling(): raise