From eebb99c049aa557d939ba0f9fe392a9bd318e51a Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Fri, 28 Aug 2026 16:58:19 +0800 Subject: [PATCH] [diffusion][kernel] avoid 4D scale-shift autotuning (#36521) --- python/sglang/kernels/ops/diffusion/README.md | 1 + .../diffusion/modulate/scale_shift_triton.py | 21 ++-- .../existing-fast-paths.md | 4 + .../diffusion/bench_scale_shift_4d.py | 114 ++++++++++++++++++ .../kernels/ops/diffusion/test_modulate.py | 29 +++++ 5 files changed, 156 insertions(+), 13 deletions(-) create mode 100644 test/registered/kernels/benchmark/diffusion/bench_scale_shift_4d.py diff --git a/python/sglang/kernels/ops/diffusion/README.md b/python/sglang/kernels/ops/diffusion/README.md index fbf78dc0c..21d74b714 100644 --- a/python/sglang/kernels/ops/diffusion/README.md +++ b/python/sglang/kernels/ops/diffusion/README.md @@ -90,6 +90,7 @@ Several norms look interchangeable and are not. Start here. | Entry point | Backend | Contract | Applies to | |---|---|---|---| +| `fuse_scale_shift_kernel` | Triton | close | contiguous BLC; scalar/row/token modulation plus causal-video `[B, F, 1, C]`, using a static capped power-of-two tile to avoid request-time autotuning | | `fused_rmsnorm_scale_shift_bitexact` | Triton | bit-exact vs flashinfer CuTe RMSNorm + aten modulate | bf16, contiguous rows, `H == 64 * threads_per_row` | | `fused_scale_residual_rmsnorm_scale_shift_bitexact` | Triton | bit-exact, incl. the preceding residual-gate add | as above | | `fused_layernorm_modulate` | Triton | bit-exact vs aten `vectorized_layer_norm` | bf16, `N % 4 == 0`, 16B-aligned | diff --git a/python/sglang/kernels/ops/diffusion/modulate/scale_shift_triton.py b/python/sglang/kernels/ops/diffusion/modulate/scale_shift_triton.py index 09b2e2a74..6741c1a2e 100644 --- a/python/sglang/kernels/ops/diffusion/modulate/scale_shift_triton.py +++ b/python/sglang/kernels/ops/diffusion/modulate/scale_shift_triton.py @@ -272,16 +272,6 @@ def _fused_residual_layernorm_scale_shift_gate_select01_kernel( tl.store(gate_row_ptr + cols, gate, mask=mask) -@triton.autotune( - configs=[ - triton.Config({"BLOCK_N": 64}, num_warps=2), - triton.Config({"BLOCK_N": 128}, num_warps=4), - triton.Config({"BLOCK_N": 256}, num_warps=4), - triton.Config({"BLOCK_N": 512}, num_warps=4), - triton.Config({"BLOCK_N": 1024}, num_warps=8), - ], - key=["inner_dim"], -) @triton.jit def _fused_scale_shift_4d_kernel( output_ptr, @@ -416,9 +406,12 @@ def fuse_scale_shift_kernel( x_2d = x.view(rows, C) output_2d = output.view(rows, C) - def grid(meta): - return (rows, triton.cdiv(C, meta["BLOCK_N"])) - + # Autotuning this bandwidth-bound kernel is much more expensive than + # the launch itself on causal video models. A capped power-of-two + # tile is fastest or within noise across the production hidden sizes. + block_n = max(64, min(512, triton.next_power_of_2(C))) + num_warps = 2 if block_n == 64 else 4 + grid = (rows, triton.cdiv(C, block_n)) num_frames = scale.shape[1] assert ( L % num_frames == 0 @@ -454,6 +447,8 @@ def fuse_scale_shift_kernel( L, num_frames, frame_seqlen, + BLOCK_N=block_n, + num_warps=num_warps, ) else: # 2D: [B, C] or [1, C] -> treat as [B, 1, C] and broadcast over L 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 98f79e895..c627f1065 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 @@ -77,6 +77,10 @@ framework-specific optimization workflow. - Locations: `elementwise.py`, `layernorm.py`, `fused_scale_shift_gate.py`, `qwen_image.py`, `triton/scale_shift.py` - Use cases: `x * (1 + scale) + shift`, `a * (k + b) + c`, and Qwen-style `(layernorm/residual layernorm) + scale/shift + gate select`. - Constraints: `x` must be CUDA and contiguous. `scale/shift` support 0D/1D/2D/3D/4D broadcast. 4D `[B, F, 1, C]` requires `L % F == 0`. +- Causal-video cold start: the 4D path uses a static capped power-of-two + column tile rather than Triton autotuning. Do not reintroduce request-time + autotuning here: LingBot-World calls this path once per transformer block, + and tuning overhead can dominate its first denoise step. - NPU fallback: `scale_shift.py` swaps to `npu_fallback` native path. - Validation: `test/registered/kernels/ops/diffusion/test_qwen_image_modulation.py`. diff --git a/test/registered/kernels/benchmark/diffusion/bench_scale_shift_4d.py b/test/registered/kernels/benchmark/diffusion/bench_scale_shift_4d.py new file mode 100644 index 000000000..dfc02a892 --- /dev/null +++ b/test/registered/kernels/benchmark/diffusion/bench_scale_shift_4d.py @@ -0,0 +1,114 @@ +import random +import sys +import time +from dataclasses import dataclass + +import torch + +from sglang.kernels.ops.diffusion import fuse_scale_shift_kernel +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.utils import is_in_ci + +register_cuda_ci( + est_time=25, stage="base-b-kernel-benchmark", runner_config="1-gpu-large" +) + + +@dataclass(frozen=True) +class Workload: + name: str + shape: tuple[int, int, int] + num_frames: int + + +FULL_WORKLOADS = [ + Workload("wan_s24960_c1536", (1, 24960, 1536), 5), + Workload("sana_video_s7800_c2240", (1, 7800, 2240), 5), + Workload("longlive_s1560_c3072", (1, 1560, 3072), 3), + Workload("lingbot_world_s4680_c5120", (1, 4680, 5120), 1), +] +CI_WORKLOADS = [ + Workload("ci_s1024_c1536", (1, 1024, 1536), 4), + Workload("ci_s512_c5120", (1, 512, 5120), 2), +] + + +def cuda_event_us(fn, warmups: int, repeats: int, rounds: int) -> float: + for _ in range(warmups): + fn() + torch.cuda.synchronize() + + samples = [] + for _ in range(rounds): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(repeats): + fn() + end.record() + end.synchronize() + samples.append(start.elapsed_time(end) * 1000.0 / repeats) + samples.sort() + return samples[len(samples) // 2] + + +def benchmark() -> None: + if not torch.cuda.is_available(): + print("CUDA required") + return + + torch.manual_seed(20260826) + random.seed(20260826) + torch.cuda.set_device(0) + + workloads = CI_WORKLOADS if is_in_ci() else FULL_WORKLOADS + warmups = 5 if is_in_ci() else 20 + repeats = 5 if is_in_ci() else 20 + rounds = 5 if is_in_ci() else 13 + + print("| workload | cold ms | torch us | triton us | speedup |") + print("|---|---:|---:|---:|---:|") + + for workload in workloads: + batch, seq_len, hidden = workload.shape + x = torch.randn(workload.shape, device="cuda", dtype=torch.bfloat16) + scale = torch.randn( + (batch, workload.num_frames, 1, hidden), + device="cuda", + dtype=torch.bfloat16, + ) + shift = torch.randn_like(x) + frame_seqlen = seq_len // workload.num_frames + + torch_fn = lambda: ( + x.unflatten(1, (workload.num_frames, frame_seqlen)) * (1 + scale) + + shift.unflatten(1, (workload.num_frames, frame_seqlen)) + ).flatten(1, 2) + triton_fn = lambda: fuse_scale_shift_kernel(x, scale, shift) + + torch.cuda.synchronize() + start = time.perf_counter() + triton_out = triton_fn() + torch.cuda.synchronize() + cold_ms = (time.perf_counter() - start) * 1000.0 + torch.testing.assert_close(triton_out, torch_fn(), atol=5e-2, rtol=5e-2) + + providers = ["torch", "triton"] + random.shuffle(providers) + fns = {"torch": torch_fn, "triton": triton_fn} + times = { + provider: cuda_event_us( + fns[provider], warmups=warmups, repeats=repeats, rounds=rounds + ) + for provider in providers + } + print( + f"| {workload.name} | {cold_ms:.2f} | {times['torch']:.2f} | " + f"{times['triton']:.2f} | {times['torch'] / times['triton']:.2f}x |" + ) + torch.cuda.empty_cache() + + +if __name__ == "__main__": + benchmark() + sys.exit(0) diff --git a/test/registered/kernels/ops/diffusion/test_modulate.py b/test/registered/kernels/ops/diffusion/test_modulate.py index eed0cfa4a..b6f4b11db 100644 --- a/test/registered/kernels/ops/diffusion/test_modulate.py +++ b/test/registered/kernels/ops/diffusion/test_modulate.py @@ -18,6 +18,7 @@ from sglang.kernels.ops.diffusion import ( can_use_residual_gate_add_cuda, fuse_layernorm_scale_shift_gate_select01_kernel, fuse_residual_layernorm_scale_shift_gate_select01_kernel, + fuse_scale_shift_kernel, ltx2_ada_values9, modulate_scale_shift, modulate_scale_shift_cuda, @@ -94,6 +95,34 @@ def test_modulate_scale_shift_guards_reject_fp32(): assert torch.equal(modulate_scale_shift(x, row, row), _eager_modulate(x, row, row)) +# Causal Wan and LingBot use per-frame 4D modulation with a per-token shift. +SCALE_SHIFT_4D_CASES = [ + ((1, 18, 96), 3), + ((2, 20, 384), 4), + ((1, 9, 1536), 3), + ((1, 4, 5120), 2), +] + + +@pytest.mark.parametrize("shape,num_frames", SCALE_SHIFT_4D_CASES) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("scale_constant", [0, 1]) +def test_scale_shift_4d_matches_torch(shape, num_frames, dtype, scale_constant): + batch, seq_len, hidden = shape + x = torch.randn(shape, device=DEVICE, dtype=dtype) + scale = torch.randn((batch, num_frames, 1, hidden), device=DEVICE, dtype=dtype) + shift = torch.randn_like(x) + + frame_seqlen = seq_len // num_frames + expected = ( + x.unflatten(1, (num_frames, frame_seqlen)) * (scale_constant + scale) + + shift.unflatten(1, (num_frames, frame_seqlen)) + ).flatten(1, 2) + actual = fuse_scale_shift_kernel(x, scale, shift, scale_constant) + + torch.testing.assert_close(actual, expected, atol=5e-2, rtol=5e-2) + + # --------------------------------------------------------------------------- # residual + gate * update # ---------------------------------------------------------------------------