[diffusion][kernel] avoid 4D scale-shift autotuning (#36521)

2026-08-28 16:58:19 +08:00
parent 45424d8434
commit eebb99c049
5 changed files with 156 additions and 13 deletions
@@ -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 |
@@ -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
@@ -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`.

@@ -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)
@@ -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
 # ---------------------------------------------------------------------------