[diffusion] feat: pick the attention backend by measuring it (#38689)

Co-authored-by: Mick Qian <mickqian@users.noreply.github.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
2026-09-12 21:52:16 +08:00
parent dc3171c322
commit 7b89b95168
6 changed files with 412 additions and 6 deletions
@@ -0,0 +1,179 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Choose a layer's attention backend by timing the candidates on its own tensors.
+
+Which backend is fastest is not a property of the GPU alone. On sm12x no FA
+kernel exists and cuDNN beats torch's flash path at head_dim 128 but loses to it
+at head_dim 64 and long sequences; on Hopper the FA backend beats cuDNN. Nor do
+synthetic timings settle it: the tensors here are non-contiguous views into a
+packed QKV buffer and backends differ in how they take that, so the measurement
+uses what the layer was actually handed, on its first forward large enough to be
+worth deciding on.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionImpl,
+    wrap_attention_impl_forward,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+# Backends only separate on the calls that carry the runtime; tuning on a short
+# text or audio stream picks the wrong winner for the long video one.
+_MIN_TUNE_NUMEL = 4 << 20
+# A candidate has to beat the incumbent by more than the spread of these timings.
+_MIN_RELATIVE_GAIN = 0.02
+_WARMUP_ITERS = 3
+_TIMED_ITERS = 8
+# A backend that disagrees this much is not computing the same attention,
+# whatever its timing says.
+_MAX_OUTPUT_DEVIATION = 0.05
+
+_reported = False
+
+
+def _timed(impl: AttentionImpl, args, kwargs) -> float:
+    for _ in range(_WARMUP_ITERS):
+        impl.forward(*args, **kwargs)
+    torch.cuda.synchronize()
+    start, end = torch.cuda.Event(True), torch.cuda.Event(True)
+    start.record()
+    for _ in range(_TIMED_ITERS):
+        impl.forward(*args, **kwargs)
+    end.record()
+    torch.cuda.synchronize()
+    return start.elapsed_time(end) / _TIMED_ITERS
+
+
+def _leading_tensor(out):
+    return out[0] if isinstance(out, (tuple, list)) else out
+
+
+def _agrees(candidate_out, reference_out) -> bool:
+    got, want = _leading_tensor(candidate_out), _leading_tensor(reference_out)
+    if not (isinstance(got, torch.Tensor) and isinstance(want, torch.Tensor)):
+        return False
+    if got.shape != want.shape:
+        return False
+    scale = want.float().abs().max().clamp_min(1e-3)
+    return bool(
+        ((got.float() - want.float()).abs().max() / scale) <= _MAX_OUTPUT_DEVIATION
+    )
+
+
+def _candidates(layer) -> list[tuple[str, AttentionImpl, AttentionBackendEnum]]:
+    from sglang.multimodal_gen.runtime.layers.attention.selector import get_attn_backend
+
+    ctor_kwargs = layer._attn_impl_ctor_kwargs
+    built: list[tuple[str, AttentionImpl, AttentionBackendEnum]] = []
+
+    # Whether SDPA is allowed to reach for cuDNN is a backend choice of its own,
+    # and on sm12x it is the only one there is.
+    if layer.backend is AttentionBackendEnum.TORCH_SDPA:
+        flipped = not ctor_kwargs.get("allow_cudnn_sdp", False)
+        built.append(
+            (
+                f"torch_sdpa(cudnn={flipped})",
+                type(layer.attn_impl)(**{**ctor_kwargs, "allow_cudnn_sdp": flipped}),
+                layer.backend,
+            )
+        )
+
+    for target in sorted(
+        layer._supported_attention_backends or (), key=lambda backend: backend.name
+    ):
+        if target is layer.backend:
+            continue
+        try:
+            backend_cls = get_attn_backend(
+                layer.head_size,
+                layer.dtype,
+                supported_attention_backends=layer._supported_attention_backends,
+                selected_attention_backend=target,
+            )
+            if backend_cls.get_enum() is not target:
+                continue
+            built.append(
+                (target.name.lower(), backend_cls.get_impl_cls()(**ctor_kwargs), target)
+            )
+        except Exception as exc:
+            logger.debug("attention autotune: %s unavailable (%s)", target, exc)
+    return built
+
+
+def _choose(layer, args, kwargs) -> tuple[AttentionImpl, AttentionBackendEnum] | None:
+    """The fastest candidate that agrees with the incumbent, or None to keep it."""
+    global _reported
+
+    incumbent = layer.attn_impl
+    reference = incumbent.forward(*args, **kwargs)
+    incumbent_label = f"{layer.backend.name.lower()} (current)"
+    timings: dict[str, tuple[float, AttentionImpl | None, AttentionBackendEnum]] = {
+        incumbent_label: (_timed(incumbent, args, kwargs), None, layer.backend)
+    }
+
+    for label, candidate, enum in _candidates(layer):
+        try:
+            output = candidate.forward(*args, **kwargs)
+        except Exception as exc:
+            logger.debug("attention autotune: %s failed (%s)", label, exc)
+            continue
+        if not _agrees(output, reference):
+            logger.debug(
+                "attention autotune: %s disagrees with %s", label, incumbent_label
+            )
+            continue
+        timings[label] = (_timed(candidate, args, kwargs), candidate, enum)
+
+    incumbent_ms = timings[incumbent_label][0]
+    best_label = min(timings, key=lambda label: timings[label][0])
+    best_ms = timings[best_label][0]
+
+    say = logger.debug if _reported else logger.info
+    _reported = True
+    report = ", ".join(f"{label} {ms:.3f}ms" for label, (ms, *_) in timings.items())
+    if best_label == incumbent_label or best_ms > incumbent_ms * (
+        1 - _MIN_RELATIVE_GAIN
+    ):
+        say("attention autotune: keeping %s (%s)", incumbent_label, report)
+        return None
+    say(
+        "attention autotune: %s -> %s, %.1f%% faster (%s)",
+        incumbent_label,
+        best_label,
+        100 * (1 - best_ms / incumbent_ms),
+        report,
+    )
+    return timings[best_label][1], timings[best_label][2]
+
+
+def install(layer) -> None:
+    """Tune this layer on its first forward worth measuring, then step aside."""
+    impl = layer.attn_impl
+    default_forward = impl.forward
+
+    def tuning_forward(*args, **kwargs):
+        query = args[0] if args else kwargs.get("query")
+        if not isinstance(query, torch.Tensor) or query.numel() < _MIN_TUNE_NUMEL:
+            return default_forward(*args, **kwargs)
+        impl.forward = default_forward
+        try:
+            winner = _choose(layer, args, kwargs)
+        except Exception as exc:
+            logger.warning_once(
+                f"attention autotune failed, keeping the default: {exc}"
+            )
+            return default_forward(*args, **kwargs)
+        if winner is None:
+            return default_forward(*args, **kwargs)
+        impl_choice, backend_choice = winner
+        layer.attn_impl = wrap_attention_impl_forward(impl_choice)
+        layer.backend = backend_choice
+        return layer.attn_impl.forward(*args, **kwargs)
+
+    impl.forward = tuning_forward
@@ -65,7 +65,7 @@ class SDPAImpl(AttentionImpl):

    def _sdpa_context(self, query: torch.Tensor):
        if self.allow_cudnn_sdp and query.device.type == "cuda":
-            return sdpa_kernel(_PYTORCH_DEFAULT_CUDA_SDP_BACKENDS)
+            return sdpa_kernel(_PYTORCH_DEFAULT_CUDA_SDP_BACKENDS, set_priority=True)
        return nullcontext()

    def forward(
@@ -413,6 +413,7 @@ class UlyssesAttention(nn.Module):
        )
        self.attn_impl = impl_cls(**self._attn_impl_ctor_kwargs)
        wrap_attention_impl_forward(self.attn_impl)
+        _maybe_install_backend_autotune(self, attn_backend.get_enum())
        self.num_heads = num_heads
        self.head_size = head_size
        self.num_kv_heads = num_kv_heads
@@ -681,6 +682,7 @@ class LocalAttention(nn.Module):
        )
        self.attn_impl = impl_cls(**self._attn_impl_ctor_kwargs)
        wrap_attention_impl_forward(self.attn_impl)
+        _maybe_install_backend_autotune(self, attn_backend.get_enum())
        self.num_heads = num_heads
        self.head_size = head_size
        self.num_kv_heads = num_kv_heads
@@ -748,7 +750,7 @@ class LocalAttention(nn.Module):
                v_ = v_.repeat_interleave(repeat_factor, dim=1)

            sdpa_context = (
-                sdpa_kernel(_PYTORCH_DEFAULT_CUDA_SDP_BACKENDS)
+                sdpa_kernel(_PYTORCH_DEFAULT_CUDA_SDP_BACKENDS, set_priority=True)
                if self.allow_cudnn_sdp and q_.device.type == "cuda"
                else nullcontext()
            )
@@ -853,6 +855,7 @@ class USPAttention(nn.Module):
        )
        self.attn_impl = impl_cls(**self._attn_impl_ctor_kwargs)
        wrap_attention_impl_forward(self.attn_impl)
+        _maybe_install_backend_autotune(self, attn_backend.get_enum())
        self.num_heads = num_heads
        self.head_size = head_size
        self.num_kv_heads = num_kv_heads
@@ -1185,7 +1188,7 @@ class USPAttention(nn.Module):
                v_ = v.transpose(1, 2)
                mask = _prepare_sdpa_mask(attn_mask, dtype=q_.dtype, device=q_.device)
                sdpa_context = (
-                    sdpa_kernel(_PYTORCH_DEFAULT_CUDA_SDP_BACKENDS)
+                    sdpa_kernel(_PYTORCH_DEFAULT_CUDA_SDP_BACKENDS, set_priority=True)
                    if self.allow_cudnn_sdp and q_.device.type == "cuda"
                    else nullcontext()
                )
@@ -1357,7 +1360,7 @@ class USPAttention(nn.Module):
            v_ = v.transpose(1, 2)
            mask = _prepare_sdpa_mask(gathered_mask, dtype=q_.dtype, device=q_.device)
            sdpa_context = (
-                sdpa_kernel(_PYTORCH_DEFAULT_CUDA_SDP_BACKENDS)
+                sdpa_kernel(_PYTORCH_DEFAULT_CUDA_SDP_BACKENDS, set_priority=True)
                if self.allow_cudnn_sdp and q_.device.type == "cuda"
                else nullcontext()
            )
@@ -1644,7 +1647,7 @@ class USPAttention(nn.Module):
            v_ = v_.repeat_interleave(repeat_factor, dim=1)

        sdpa_context = (
-            sdpa_kernel(_PYTORCH_DEFAULT_CUDA_SDP_BACKENDS)
+            sdpa_kernel(_PYTORCH_DEFAULT_CUDA_SDP_BACKENDS, set_priority=True)
            if self.allow_cudnn_sdp and q_.device.type == "cuda"
            else nullcontext()
        )
@@ -1863,7 +1866,7 @@ class USPAttention(nn.Module):
        v_ = v.transpose(1, 2)
        mask = _prepare_sdpa_mask(attn_mask, dtype=q_.dtype, device=q_.device)
        sdpa_context = (
-            sdpa_kernel(_PYTORCH_DEFAULT_CUDA_SDP_BACKENDS)
+            sdpa_kernel(_PYTORCH_DEFAULT_CUDA_SDP_BACKENDS, set_priority=True)
            if self.allow_cudnn_sdp and q_.device.type == "cuda"
            else nullcontext()
        )
@@ -2097,3 +2100,21 @@ for _attn_cls in (
 ):
    _attn_cls.forward = _make_breakable_attention_forward(_attn_cls.forward)
 del _attn_cls
+
+
+def _maybe_install_backend_autotune(layer, backend) -> None:
+    """Opt-in: let the layer pick its backend by measurement on its first big call."""
+    from sglang.multimodal_gen.runtime.server_args import get_global_server_args
+
+    try:
+        if not get_global_server_args().enable_attention_backend_autotune:
+            return
+    except Exception:  # no ServerArgs yet (unit tests, tooling)
+        return
+    if getattr(layer, "_required_attention_backend", None) is not None:
+        return
+    from sglang.multimodal_gen.runtime.layers.attention.autotune import install
+
+    layer.backend = backend
+    layer._default_attn_backend = backend
+    install(layer)
@@ -275,6 +275,10 @@ class ServerArgs(DisaggServerArgsMixin):

    # Attention
    attention_backend: str = None
+    # Time the viable attention backends on the model's own tensors during
+    # warmup and keep the fastest per layer. Off by default: it has only been
+    # measured on sm90 and sm12x.
+    enable_attention_backend_autotune: bool = False
    attention_backend_config: addict.Dict | None = None
    component_attention_backends: dict[str, str] | str | None = field(
        default_factory=dict
@@ -0,0 +1,108 @@
+# SPDX-License-Identifier: Apache-2.0
+"""The rules the autotuner has to obey: only switch on a clear, correct win.
+
+Timing is stubbed here so the rules are what is under test, not the GPU.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from sglang.multimodal_gen.runtime.layers.attention import autotune
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+
+QUERY = torch.zeros(4, 4)
+REFERENCE = torch.ones(2, 2)
+
+
+class _Impl:
+    def __init__(self, output=REFERENCE):
+        self._output = output
+        self.calls = 0
+
+    def forward(self, *args, **kwargs):
+        self.calls += 1
+        return self._output
+
+
+def _layer(incumbent):
+    return SimpleNamespace(
+        attn_impl=incumbent,
+        backend=AttentionBackendEnum.TORCH_SDPA,
+        head_size=128,
+        dtype=torch.bfloat16,
+        _attn_impl_ctor_kwargs={},
+        _supported_attention_backends=set(),
+    )
+
+
+@pytest.fixture
+def stub(monkeypatch):
+    """Drive _choose with fixed candidates and fixed timings."""
+
+    def install(candidates: list[tuple[str, _Impl, object]], timings: dict[int, float]):
+        monkeypatch.setattr(autotune, "_candidates", lambda layer: candidates)
+        monkeypatch.setattr(
+            autotune, "_timed", lambda impl, args, kwargs: timings[id(impl)]
+        )
+
+    return install
+
+
+def test_keeps_the_incumbent_without_a_clear_win(stub):
+    incumbent, rival = _Impl(), _Impl()
+    stub(
+        [("rival", rival, AttentionBackendEnum.FA)],
+        {id(incumbent): 10.0, id(rival): 9.9},  # 1%, under the margin
+    )
+    assert autotune._choose(_layer(incumbent), (QUERY,), {}) is None
+
+
+def test_switches_when_a_candidate_wins_by_more_than_the_margin(stub):
+    incumbent, rival = _Impl(), _Impl()
+    stub(
+        [("rival", rival, AttentionBackendEnum.FA)],
+        {id(incumbent): 10.0, id(rival): 8.0},
+    )
+    chosen = autotune._choose(_layer(incumbent), (QUERY,), {})
+    assert chosen is not None
+    assert chosen[0] is rival
+    assert chosen[1] is AttentionBackendEnum.FA
+
+
+def test_a_faster_candidate_that_disagrees_is_rejected(stub):
+    incumbent = _Impl()
+    wrong = _Impl(output=REFERENCE * 5)
+    stub(
+        [("wrong", wrong, AttentionBackendEnum.FA)],
+        {id(incumbent): 10.0, id(wrong): 1.0},
+    )
+    assert autotune._choose(_layer(incumbent), (QUERY,), {}) is None
+
+
+def test_a_candidate_that_raises_is_skipped(stub, monkeypatch):
+    incumbent = _Impl()
+    broken = _Impl()
+
+    def explode(*args, **kwargs):
+        raise RuntimeError("unsupported here")
+
+    broken.forward = explode
+    stub([("broken", broken, AttentionBackendEnum.FA)], {id(incumbent): 10.0})
+    assert autotune._choose(_layer(incumbent), (QUERY,), {}) is None
+
+
+def test_small_calls_stay_on_the_default_and_leave_the_tuner_armed(monkeypatch):
+    incumbent = _Impl()
+    layer = _layer(incumbent)
+    called = []
+    monkeypatch.setattr(autotune, "_choose", lambda *a, **k: called.append(1))
+
+    autotune.install(layer)
+    small = torch.zeros(8, 8)
+    assert small.numel() < autotune._MIN_TUNE_NUMEL
+    incumbent.forward(small)
+
+    assert called == [], "tuning must wait for a call worth measuring"
+    assert layer.attn_impl is incumbent
@@ -0,0 +1,94 @@
+# SPDX-License-Identifier: Apache-2.0
+"""``allow_cudnn_sdp`` has to override torch's own backend choice.
+
+The backend list is written cuDNN-first, but ``sdpa_kernel`` treats it as an
+allow-set unless ``set_priority`` is passed -- and it is the same set torch
+already chooses from, so without that flag the context is inert. Only the kernel
+that actually ran distinguishes the two, so that is what these assert on.
+
+Where the allow-set alone already lands on cuDNN (Hopper, for one) the flag has
+nothing left to do and the check skips rather than asserting something it cannot
+observe.
+"""
+
+from contextlib import nullcontext
+
+import pytest
+import torch
+import torch.nn.functional as F
+from torch.nn.attention import SDPBackend, sdpa_kernel
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.sdpa import (
+    _PYTORCH_DEFAULT_CUDA_SDP_BACKENDS,
+    SDPAImpl,
+)
+
+NUM_HEADS, HEAD_DIM, SEQ_LEN = 8, 128, 512
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="attention dispatch needs CUDA"
+)
+
+
+def _qkv():
+    shape = (1, NUM_HEADS, SEQ_LEN, HEAD_DIM)
+    return tuple(
+        torch.randn(shape, device="cuda", dtype=torch.bfloat16) for _ in range(3)
+    )
+
+
+def _kernel_under(context, q, k, v) -> str:
+    """Which family of attention kernel ran inside ``context``."""
+    from torch.profiler import ProfilerActivity, profile
+
+    with profile(activities=[ProfilerActivity.CUDA]) as prof:
+        with context:
+            F.scaled_dot_product_attention(q, k, v)
+        torch.cuda.synchronize()
+    for event in prof.events():
+        if event.device_type.name != "CUDA" or not event.self_device_time_total:
+            continue
+        name = event.key.lower()
+        # cuDNN names its own kernels `cudnn_generated_..._flash_...`, so cuDNN
+        # has to be checked first.
+        if "cudnn" in name:
+            return "cudnn"
+        if "flash" in name or "fmha" in name:
+            return "flash"
+    return "unknown"
+
+
+def _impl(allow_cudnn_sdp: bool) -> SDPAImpl:
+    return SDPAImpl(
+        num_heads=NUM_HEADS,
+        head_size=HEAD_DIM,
+        causal=False,
+        softmax_scale=HEAD_DIM**-0.5,
+        allow_cudnn_sdp=allow_cudnn_sdp,
+    )
+
+
+def test_allow_cudnn_sdp_beats_the_backend_torch_would_pick():
+    q, k, v = _qkv()
+    try:
+        with sdpa_kernel(SDPBackend.CUDNN_ATTENTION):
+            F.scaled_dot_product_attention(q, k, v)
+    except RuntimeError:
+        pytest.skip("no cuDNN attention kernel for this shape on this GPU")
+
+    # The allow-set on its own is what this code did before priority was passed.
+    without_priority = sdpa_kernel(_PYTORCH_DEFAULT_CUDA_SDP_BACKENDS)
+    if _kernel_under(without_priority, q, k, v) == "cudnn":
+        pytest.skip("the allow-set alone already lands on cuDNN on this GPU")
+
+    assert _kernel_under(_impl(True)._sdpa_context(q), q, k, v) == "cudnn"
+
+
+def test_opting_out_leaves_backend_selection_alone():
+    q, _, _ = _qkv()
+    assert isinstance(_impl(False)._sdpa_context(q), type(nullcontext()))
+
+
+def test_cpu_tensors_do_not_get_a_cuda_context():
+    cpu_q = torch.randn(1, NUM_HEADS, 16, HEAD_DIM)
+    assert isinstance(_impl(True)._sdpa_context(cpu_q), type(nullcontext()))