[diffusion][kernel] support transposed residual-gate add (#36504)

python/sglang/kernels/jit/csrc/diffusion/residual_gate_add.cuh

@@ -25,6 +25,8 @@ constexpr uint32_t kBroadcastRowsPerBlock = 4;
 constexpr uint32_t kBroadcastColsPerBlock = 256;
 constexpr uint32_t kMaxGrid = 65535;
 constexpr uintptr_t kAlignment = 16;
+constexpr uint32_t kTransposeTile = 32;
+constexpr uint32_t kTransposeBlockSize = 256;
 
 enum class GateMode : int { kFull, kBroadcastRow };
 
@@ -110,6 +112,55 @@ __global__ void residual_gate_add_scalar_kernel(
   }
 }
 
+/**
+ * Fuse the residual-gate update when the residual/output use the common
+ * ``[batch, hidden, tokens]`` backing layout exposed as ``[batch, tokens,
+ * hidden]``, while update remains contiguous in the logical layout.
+ *
+ * Loading update in logical row-major order and consuming it transposed from
+ * shared memory keeps both update reads and residual/output traffic coalesced.
+ */
+template <typename T>
+__global__ void residual_gate_add_transposed_kernel(
+    T* __restrict__ out,
+    const T* __restrict__ residual,
+    const T* __restrict__ update,
+    const T* __restrict__ gate,
+    int64_t tokens,
+    int64_t hidden_size) {
+  __shared__ T update_tile[kTransposeTile][kTransposeTile + 1];
+
+  const int64_t batch = blockIdx.z;
+  const int64_t token_base = static_cast<int64_t>(blockIdx.x) * kTransposeTile;
+  const int64_t hidden_base = static_cast<int64_t>(blockIdx.y) * kTransposeTile;
+  const int64_t batch_offset = batch * tokens * hidden_size;
+
+#pragma unroll
+  for (uint32_t item = threadIdx.x; item < kTransposeTile * kTransposeTile; item += kTransposeBlockSize) {
+    const uint32_t token_in_tile = item / kTransposeTile;
+    const uint32_t hidden_in_tile = item % kTransposeTile;
+    const int64_t token = token_base + token_in_tile;
+    const int64_t hidden = hidden_base + hidden_in_tile;
+    if (token < tokens && hidden < hidden_size) {
+      update_tile[token_in_tile][hidden_in_tile] = update[batch_offset + token * hidden_size + hidden];
+    }
+  }
+  __syncthreads();
+
+#pragma unroll
+  for (uint32_t item = threadIdx.x; item < kTransposeTile * kTransposeTile; item += kTransposeBlockSize) {
+    const uint32_t hidden_in_tile = item / kTransposeTile;
+    const uint32_t token_in_tile = item % kTransposeTile;
+    const int64_t token = token_base + token_in_tile;
+    const int64_t hidden = hidden_base + hidden_in_tile;
+    if (token < tokens && hidden < hidden_size) {
+      const int64_t transposed_offset = batch_offset + hidden * tokens + token;
+      out[transposed_offset] = residual_gate_value(
+          residual[transposed_offset], update_tile[token_in_tile][hidden_in_tile], SGLANG_LDG(gate + hidden));
+    }
+  }
+}
+
 }  // namespace
 
 /**
@@ -201,6 +252,49 @@ struct ResidualGateAddKernel {
           hidden_size);
     }
   }
+
+  static void run_transposed(
+      tvm::ffi::TensorView out, tvm::ffi::TensorView residual, tvm::ffi::TensorView update, tvm::ffi::TensorView gate) {
+    using namespace host;
+
+    auto B = SymbolicSize{"batch"};
+    auto S = SymbolicSize{"tokens"};
+    auto D = SymbolicSize{"hidden_size"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+    TensorMatcher({B, S, D}).with_strides({-1, 1, -1}).with_dtype<T>().with_device(device).verify(out).verify(residual);
+    TensorMatcher({B, S, D}).with_strides({-1, -1, 1}).with_dtype<T>().with_device(device).verify(update);
+    TensorMatcher({1, 1, D}).with_strides({-1, -1, 1}).with_dtype<T>().with_device(device).verify(gate);
+
+    const int64_t batch = B.unwrap();
+    const int64_t tokens = S.unwrap();
+    const int64_t hidden_size = D.unwrap();
+    CHECK_HOST(batch > 0 && tokens > 0 && hidden_size > 0) << "transposed residual-gate tensors must be non-empty";
+    CHECK_HOST(
+        batch <= kMaxGrid && div_ceil(tokens, int64_t{kTransposeTile}) <= kMaxGrid &&
+        div_ceil(hidden_size, int64_t{kTransposeTile}) <= kMaxGrid)
+        << "transposed residual-gate grid exceeds CUDA limits";
+    const int64_t batch_stride = tokens * hidden_size;
+    CHECK_HOST(
+        out.stride(0) == batch_stride && out.stride(2) == tokens && residual.stride(0) == batch_stride &&
+        residual.stride(2) == tokens)
+        << "residual/output must use the transposed dense layout";
+    CHECK_HOST(update.stride(0) == batch_stride && update.stride(1) == hidden_size) << "update must be contiguous";
+
+    auto* out_ptr = static_cast<T*>(out.data_ptr());
+    const auto* residual_ptr = static_cast<const T*>(residual.data_ptr());
+    const auto* update_ptr = static_cast<const T*>(update.data_ptr());
+    const auto* gate_ptr = static_cast<const T*>(gate.data_ptr());
+    CHECK_HOST(out_ptr != residual_ptr && out_ptr != update_ptr && out_ptr != gate_ptr)
+        << "output must not alias an input";
+
+    const dim3 grid(
+        static_cast<uint32_t>(div_ceil(tokens, int64_t{kTransposeTile})),
+        static_cast<uint32_t>(div_ceil(hidden_size, int64_t{kTransposeTile})),
+        static_cast<uint32_t>(batch));
+    LaunchKernel(grid, kTransposeBlockSize, device.unwrap())(
+        residual_gate_add_transposed_kernel<T>, out_ptr, residual_ptr, update_ptr, gate_ptr, tokens, hidden_size);
+  }
 };
 
 }  // namespace residual_gate_add

python/sglang/kernels/ops/diffusion/README.md

@@ -110,6 +110,18 @@ Several norms look interchangeable and are not. Start here.
 | `fused_qk_head_layernorm` | Triton | bit-exact | per-head LN on q/k, `dim_head % 4 == 0`, `<= 128` |
 | `triton_one_pass_rms_norm` | Triton | close | standalone RMSNorm, one pass |
 
+### Residual gating
+
+| Entry point | Backend | Contract | Applies to |
+|---|---|---|---|
+| `residual_gate_add` | JIT CUDA | bit-exact `residual + update * gate` | contiguous tensors, or a transposed-dense `[B, tokens, hidden]` residual/output with contiguous update and row-broadcast gate (SANA-Video) |
+
+The transposed-dense path uses a shared-memory tile to read the update in
+logical row-major order while keeping residual reads and output writes
+coalesced in their `[B, hidden, tokens]` backing layout. Do not insert a
+`.contiguous()` merely to reach the ordinary path; that restores an entire
+tensor copy per residual site.
+
 ### RoPE / QK-norm
 
 | Entry point | Backend | Contract |

python/sglang/kernels/ops/diffusion/modulate/residual_gate_add_jit.py

@@ -14,6 +14,8 @@ if TYPE_CHECKING:
 
 _SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32)
 _BIT_EXACT_DTYPES = (torch.float16, torch.bfloat16)
+_TRANSPOSE_TILE = 32
+_MAX_GRID_DIM = 65535
 _FAILED_RUNTIME_KEYS: set[tuple[int | None, torch.dtype]] = set()
 
 logger = logging.getLogger(__name__)
@@ -33,6 +35,10 @@ def _jit_residual_gate_add_module(dtype: torch.dtype) -> Module:
                 "residual_gate_add",
                 "residual_gate_add::" f"ResidualGateAddKernel<{args}>::run",
             ),
+            (
+                "residual_gate_add_transposed",
+                "residual_gate_add::" f"ResidualGateAddKernel<{args}>::run_transposed",
+            ),
         ],
     )
 
@@ -40,7 +46,12 @@ def _jit_residual_gate_add_module(dtype: torch.dtype) -> Module:
 def _fake_impl(
     residual: torch.Tensor, update: torch.Tensor, gate: torch.Tensor
 ) -> torch.Tensor:
-    return torch.empty_like(residual)
+    return torch.empty_strided(
+        residual.shape,
+        residual.stride(),
+        dtype=residual.dtype,
+        device=residual.device,
+    )
 
 
 @register_custom_op(
@@ -51,8 +62,16 @@ def _fake_impl(
 def _residual_gate_add_custom_op(
     residual: torch.Tensor, update: torch.Tensor, gate: torch.Tensor
 ) -> torch.Tensor:
-    out = torch.empty_like(residual)
+    out = torch.empty_strided(
+        residual.shape,
+        residual.stride(),
+        dtype=residual.dtype,
+        device=residual.device,
+    )
     module = _jit_residual_gate_add_module(residual.dtype)
+    if _is_transposed_dense_residual(residual, update, gate):
+        module.residual_gate_add_transposed(out, residual, update, gate)
+        return out
     broadcast_gate = gate.shape != residual.shape
     module.residual_gate_add(
         out.view(-1),
@@ -71,6 +90,22 @@ def _is_row_broadcast_gate(residual: torch.Tensor, gate: torch.Tensor) -> bool:
     return all(size == 1 for size in gate.shape[:-1])
 
 
+def _is_transposed_dense_residual(
+    residual: torch.Tensor, update: torch.Tensor, gate: torch.Tensor
+) -> bool:
+    if residual.dim() != 3 or gate.shape != (1, 1, residual.shape[-1]):
+        return False
+    batch, tokens, hidden_size = residual.shape
+    return (
+        batch <= _MAX_GRID_DIM
+        and (tokens + _TRANSPOSE_TILE - 1) // _TRANSPOSE_TILE <= _MAX_GRID_DIM
+        and (hidden_size + _TRANSPOSE_TILE - 1) // _TRANSPOSE_TILE <= _MAX_GRID_DIM
+        and residual.stride() == (tokens * hidden_size, 1, tokens)
+        and update.is_contiguous()
+        and gate.is_contiguous()
+    )
+
+
 def can_use_residual_gate_add_cuda(
     residual: torch.Tensor, update: torch.Tensor, gate: torch.Tensor
 ) -> bool:
@@ -86,8 +121,10 @@ def can_use_residual_gate_add_cuda(
         and residual.numel() > 0
         and update.shape == residual.shape
         and (gate.shape == residual.shape or _is_row_broadcast_gate(residual, gate))
-        and residual.is_contiguous()
-        and update.is_contiguous()
+        and (
+            (residual.is_contiguous() and update.is_contiguous())
+            or _is_transposed_dense_residual(residual, update, gate)
+        )
         and gate.is_contiguous()
     )
 

python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md

@@ -171,15 +171,15 @@ framework-specific optimization workflow.
 - Constraints: `cos` and `sin` shapes must match `[B, H, S, head_dim / 2]`, and `inner_dim == H * head_dim`.
 - Workflow rule: if LTX-2 traces show a large split-RoPE PyTorch chain, check whether the LTX2-specific Triton path was disabled by shape or dtype before proposing a new RoPE kernel.
 
-10. LTX2 and LongCat-Image residual-gate add fusion
+10. Shared residual-gate add fusion (LTX2, LongCat-Image, SANA, and SANA-Video)
 - Kernel: `diffusion_residual_gate_add`
-- Locations: `kernels/ops/diffusion/modulate/residual_gate_add_jit.py`, `kernels/jit/csrc/diffusion/residual_gate_add.cuh`, `runtime/models/dits/ltx_2.py`, `runtime/models/dits/longcat_image.py`
-- Use case: `residual + update * gate` in LTX2 self-attention, prompt cross-attention, audio/video cross-attention, and feed-forward residual updates, plus LongCat-Image joint- and single-stream transformer residuals.
-- Constraints: `residual`, `update`, and `gate` must be CUDA tensors on the same device, contiguous, same dtype (`fp16`, `bf16`, or `fp32`), with `update.shape == residual.shape`; `gate` can match `residual` or be row-broadcast with the last dimension matching.
-- Behavior: LTX2 and LongCat-Image call `residual_gate_add(...)` from the kernels package directly. The CUDA custom op is used while guards pass. On a runtime exception outside `torch.compile`, it logs once, disables the fast path for the process, and falls back to `residual + update * gate`.
+- Locations: `kernels/ops/diffusion/modulate/residual_gate_add_jit.py`, `kernels/jit/csrc/diffusion/residual_gate_add.cuh`, `runtime/models/dits/ltx_2.py`, `runtime/models/dits/longcat_image.py`, `runtime/models/dits/sana.py`, and `runtime/models/dits/sana_video.py`.
+- Use case: `residual + update * gate` in LTX2 attention/MLP residuals, LongCat-Image joint- and single-stream transformer residuals, and SANA/SANA-Video transformer blocks.
+- Constraints: inputs must be same-device CUDA tensors with one dtype (`fp16`, `bf16`, or `fp32`) and `update.shape == residual.shape`. The ordinary path accepts contiguous inputs and a full or row-broadcast gate. The SANA-Video path also accepts a transposed-dense 3D residual (`stride == (tokens * hidden, 1, tokens)`), a contiguous update, and a contiguous `[1, 1, hidden]` gate; it preserves the residual stride in its output.
+- Behavior: model code calls `residual_gate_add(...)` directly. The CUDA custom op is used while guards pass. On a runtime exception outside `torch.compile`, it logs once, disables the fast path for that device/dtype, and falls back to `residual + update * gate`.
 - Validation: `test/registered/kernels/ops/diffusion/test_modulate.py`, `python/sglang/multimodal_gen/test/unit/test_longcat_image_residual_gate.py`.
 - Microbench: `test/registered/kernels/benchmark/diffusion/bench_residual_gate_add.py`.
-- Workflow rule: if LTX2 or LongCat-Image traces show repeated elementwise `mul` + `add` ladders around attention or MLP residuals, check whether this existing CUDA path was disabled by shape, dtype, contiguity, or a prior runtime failure before proposing another elementwise fusion.
+- Workflow rule: if LTX2, LongCat-Image, or SANA traces show repeated elementwise `mul` + `add` ladders around attention or MLP residuals, inspect input strides and check whether this existing CUDA path was disabled by shape, dtype, layout, or a prior runtime failure before proposing another elementwise fusion. For a transposed residual plus contiguous update, do not force `.contiguous()`; the tiled path is designed to fuse the mixed-layout access.
 
 11. MiniMax-H3 indexed AdaLN modulation and gated residual fusion
 - Kernels: `indexed_scale_shift_bf16_`, `indexed_gate_bf16_`
@@ -396,9 +396,10 @@ framework-specific optimization workflow.
   `QualityGatedFusion`, not a first-sight `BitExactFusionGate` — the fused
   kernel is <=1 ULP off aten, so it is request-gated instead of verified),
   wired at the six `LTX2TransformerBlock` adaLN sites in `ltx_2.py`.
-- LTX2 residual-gate add: `ltx_2.py` calls `residual_gate_add` from
+- Shared residual-gate add: `ltx_2.py`, `sana.py`, and `sana_video.py` call `residual_gate_add` from
   `kernels/ops/diffusion/modulate/residual_gate_add_jit.py` directly for attention,
-  cross-attention, and MLP residual updates.
+  cross-attention, and MLP residual updates; SANA-Video's transposed residual
+  uses the mixed-layout tiled kernel without an intermediate contiguous copy.
 - Wan causal VAE: `cat_pad_channels_last_3d` and `dup_up3d_add` in
   `wanvae.py`, backed by `triton/wan_causal_cache.py`.
 - Varlen USP attention: `fused_pack_qkv` and `fused_scatter_to_padded` in `attention/layer.py`.