diff --git a/python/sglang/kernels/jit/csrc/diffusion/residual_gate_add.cuh b/python/sglang/kernels/jit/csrc/diffusion/residual_gate_add.cuh index 75670e5fd..bdd3830d3 100644 --- a/python/sglang/kernels/jit/csrc/diffusion/residual_gate_add.cuh +++ b/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 +__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(blockIdx.x) * kTransposeTile; + const int64_t hidden_base = static_cast(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(); + TensorMatcher({B, S, D}).with_strides({-1, 1, -1}).with_dtype().with_device(device).verify(out).verify(residual); + TensorMatcher({B, S, D}).with_strides({-1, -1, 1}).with_dtype().with_device(device).verify(update); + TensorMatcher({1, 1, D}).with_strides({-1, -1, 1}).with_dtype().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(out.data_ptr()); + const auto* residual_ptr = static_cast(residual.data_ptr()); + const auto* update_ptr = static_cast(update.data_ptr()); + const auto* gate_ptr = static_cast(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(div_ceil(tokens, int64_t{kTransposeTile})), + static_cast(div_ceil(hidden_size, int64_t{kTransposeTile})), + static_cast(batch)); + LaunchKernel(grid, kTransposeBlockSize, device.unwrap())( + residual_gate_add_transposed_kernel, out_ptr, residual_ptr, update_ptr, gate_ptr, tokens, hidden_size); + } }; } // namespace residual_gate_add diff --git a/python/sglang/kernels/ops/diffusion/README.md b/python/sglang/kernels/ops/diffusion/README.md index 12e4d44f9..fbf78dc0c 100644 --- a/python/sglang/kernels/ops/diffusion/README.md +++ b/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 | diff --git a/python/sglang/kernels/ops/diffusion/modulate/residual_gate_add_jit.py b/python/sglang/kernels/ops/diffusion/modulate/residual_gate_add_jit.py index 6bd81aa62..1e0202aff 100644 --- a/python/sglang/kernels/ops/diffusion/modulate/residual_gate_add_jit.py +++ b/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() ) diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md index 94fbe4ac5..98f79e895 100644 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md +++ b/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`. diff --git a/test/registered/kernels/benchmark/diffusion/bench_residual_gate_add.py b/test/registered/kernels/benchmark/diffusion/bench_residual_gate_add.py index eb9d796de..9cd3f8674 100644 --- a/test/registered/kernels/benchmark/diffusion/bench_residual_gate_add.py +++ b/test/registered/kernels/benchmark/diffusion/bench_residual_gate_add.py @@ -18,6 +18,7 @@ class Workload: name: str residual_shape: tuple[int, ...] gate_shape: tuple[int, ...] + transposed_residual: bool = False FULL_WORKLOADS = [ @@ -28,10 +29,22 @@ FULL_WORKLOADS = [ Workload("flux2_bcast_s4096_c3072", (1, 4096, 3072), (1, 1, 3072)), Workload("flux2_bcast_s512_c3072", (1, 512, 3072), (1, 1, 3072)), Workload("ltx2_full_s126_c2048", (1, 126, 2048), (1, 126, 2048)), + Workload( + "sana_video_bcast_s7800_c2240_transposed", + (1, 7800, 2240), + (1, 1, 2240), + transposed_residual=True, + ), ] CI_WORKLOADS = [ Workload("ltx2_bcast_s1024_c4096", (1, 1024, 4096), (1, 1, 4096)), Workload("ltx2_full_s512_c4096", (1, 512, 4096), (1, 512, 4096)), + Workload( + "sana_video_bcast_s512_c2240_transposed", + (1, 512, 2240), + (1, 1, 2240), + transposed_residual=True, + ), ] @@ -68,31 +81,44 @@ def benchmark() -> None: repeats = 5 if is_in_ci() else 20 rounds = 5 if is_in_ci() else 13 - print("| workload | gate | torch us | triton us | cuda us | cuda/triton |") - print("|---|---|---:|---:|---:|---:|") + print( + "| workload | gate | torch us | triton us | cuda us | reference | " "ref/cuda |" + ) + print("|---|---|---:|---:|---:|---|---:|") for workload in workloads: - residual = torch.randn( + if workload.transposed_residual: + batch, tokens, hidden_size = workload.residual_shape + residual = torch.randn( + (batch, hidden_size, tokens), device="cuda", dtype=torch.bfloat16 + ).transpose(1, 2) + else: + residual = torch.randn( + workload.residual_shape, device="cuda", dtype=torch.bfloat16 + ) + update = torch.randn( workload.residual_shape, device="cuda", dtype=torch.bfloat16 ) - update = torch.randn_like(residual) gate = torch.randn(workload.gate_shape, device="cuda", dtype=torch.bfloat16) ref = residual + update * gate - triton_out = fuse_scale_shift_kernel(update, gate, residual, scale_constant=0) cuda_out = residual_gate_add_cuda(residual, update, gate) torch.cuda.synchronize() - torch.testing.assert_close(triton_out, ref, atol=5e-2, rtol=5e-2) torch.testing.assert_close(cuda_out, ref, atol=5e-2, rtol=5e-2) fns = { "torch": lambda: residual + update * gate, - "triton": lambda: fuse_scale_shift_kernel( - update, gate, residual, scale_constant=0 - ), "cuda": lambda: residual_gate_add_cuda(residual, update, gate), } - order = ["torch", "triton", "cuda"] + if not workload.transposed_residual: + triton_out = fuse_scale_shift_kernel( + update, gate, residual, scale_constant=0 + ) + torch.testing.assert_close(triton_out, ref, atol=5e-2, rtol=5e-2) + fns["triton"] = lambda: fuse_scale_shift_kernel( + update, gate, residual, scale_constant=0 + ) + order = list(fns) random.shuffle(order) times = { name: cuda_event_us(fns[name], warmups, repeats, rounds) for name in order @@ -101,10 +127,12 @@ def benchmark() -> None: gate_kind = ( "bcast" if workload.gate_shape != workload.residual_shape else "full" ) + reference = "torch" if workload.transposed_residual else "triton" + triton_us = f"{times['triton']:.2f}" if "triton" in times else "n/a" print( f"| {workload.name} | {gate_kind} | {times['torch']:.2f} | " - f"{times['triton']:.2f} | {times['cuda']:.2f} | " - f"{times['triton'] / times['cuda']:.3f}x |" + f"{triton_us} | {times['cuda']:.2f} | {reference} | " + f"{times[reference] / times['cuda']:.3f}x |" ) torch.cuda.empty_cache() diff --git a/test/registered/kernels/ops/diffusion/test_modulate.py b/test/registered/kernels/ops/diffusion/test_modulate.py index 3ed13deeb..eed0cfa4a 100644 --- a/test/registered/kernels/ops/diffusion/test_modulate.py +++ b/test/registered/kernels/ops/diffusion/test_modulate.py @@ -146,6 +146,79 @@ def test_residual_gate_add_dtypes(dtype, gate_shape): ) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("shape", [(1, 17, 65), (1, 7800, 2240), (2, 33, 128)]) +def test_residual_gate_add_transposed_residual(dtype, shape): + batch, tokens, hidden_size = shape + residual = torch.randn( + (batch, hidden_size, tokens), device=DEVICE, dtype=dtype + ).transpose(1, 2) + update = torch.randn(shape, device=DEVICE, dtype=dtype) + gate = torch.randn((1, 1, hidden_size), device=DEVICE, dtype=dtype) + + assert not residual.is_contiguous() + assert can_use_residual_gate_add_cuda(residual, update, gate) + ref = residual + update * gate + out = residual_gate_add_cuda(residual, update, gate) + _assert_gate_add(out, ref) + assert out.stride() == ref.stride() == residual.stride() + + +def test_residual_gate_add_transposed_storage_offsets(): + tokens, hidden_size = 33, 128 + residual = ( + torch.randn(1 + tokens * hidden_size, device=DEVICE, dtype=torch.bfloat16)[1:] + .view(1, hidden_size, tokens) + .transpose(1, 2) + ) + update = torch.randn(1 + tokens * hidden_size, device=DEVICE, dtype=torch.bfloat16)[ + 1: + ].view(1, tokens, hidden_size) + gate = torch.randn(1 + hidden_size, device=DEVICE, dtype=torch.bfloat16)[1:].view( + 1, 1, hidden_size + ) + + assert residual.storage_offset() > 0 + assert update.storage_offset() > 0 + assert gate.storage_offset() > 0 + assert can_use_residual_gate_add_cuda(residual, update, gate) + out = residual_gate_add_cuda(residual, update, gate) + assert torch.equal(out, residual + update * gate) + + +def test_residual_gate_add_transposed_torch_compile_fullgraph(): + residual = torch.randn((1, 128, 32), device=DEVICE, dtype=torch.bfloat16).transpose( + 1, 2 + ) + update = torch.randn_like(residual, memory_format=torch.contiguous_format) + gate = torch.randn((1, 1, 128), device=DEVICE, dtype=torch.bfloat16) + compiled = torch.compile(residual_gate_add, fullgraph=True) + out = compiled(residual, update, gate) + assert torch.equal(out, residual + update * gate) + assert out.stride() == residual.stride() + + +def test_residual_gate_add_transposed_cuda_graph(): + residual = torch.randn((1, 128, 32), device=DEVICE, dtype=torch.bfloat16).transpose( + 1, 2 + ) + update = torch.randn_like(residual, memory_format=torch.contiguous_format) + gate = torch.randn((1, 1, 128), device=DEVICE, dtype=torch.bfloat16) + + # Build the JIT module before capture; graph capture must contain only the + # allocation and kernel launch used during steady-state replay. + residual_gate_add_cuda(residual, update, gate) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = residual_gate_add_cuda(residual, update, gate) + graph.replay() + torch.cuda.synchronize() + + assert torch.equal(out, residual + update * gate) + assert out.stride() == residual.stride() + + def test_residual_gate_add_guards_and_eager_fallback(): residual = torch.randn((1, 8, 64), device=DEVICE, dtype=torch.bfloat16) update = torch.randn_like(residual)