From 7b89b95168aad38266207fc442ea50b0d15aed7d Mon Sep 17 00:00:00 2001 From: Mick Date: Sat, 12 Sep 2026 21:52:16 +0800 Subject: [PATCH] [diffusion] feat: pick the attention backend by measuring it (#38689) Co-authored-by: Mick Qian Co-authored-by: Claude Opus 5 --- .../runtime/layers/attention/autotune.py | 179 ++++++++++++++++++ .../runtime/layers/attention/backends/sdpa.py | 2 +- .../runtime/layers/attention/layer.py | 31 ++- .../runtime/server_args/server_args.py | 4 + .../unit/test_attention_backend_autotune.py | 108 +++++++++++ .../test/unit/test_sdpa_cudnn_priority.py | 94 +++++++++ 6 files changed, 412 insertions(+), 6 deletions(-) create mode 100644 python/sglang/multimodal_gen/runtime/layers/attention/autotune.py create mode 100644 python/sglang/multimodal_gen/test/unit/test_attention_backend_autotune.py create mode 100644 python/sglang/multimodal_gen/test/unit/test_sdpa_cudnn_priority.py diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/autotune.py b/python/sglang/multimodal_gen/runtime/layers/attention/autotune.py new file mode 100644 index 000000000..576de062a --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/layers/attention/autotune.py @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py index 73f55158c..5552cfa88 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py @@ -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( diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py index 4b289356a..6929d87f8 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py @@ -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) diff --git a/python/sglang/multimodal_gen/runtime/server_args/server_args.py b/python/sglang/multimodal_gen/runtime/server_args/server_args.py index b3c00ca5a..22f0ec22a 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args/server_args.py @@ -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 diff --git a/python/sglang/multimodal_gen/test/unit/test_attention_backend_autotune.py b/python/sglang/multimodal_gen/test/unit/test_attention_backend_autotune.py new file mode 100644 index 000000000..3f29fbb3b --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_attention_backend_autotune.py @@ -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 diff --git a/python/sglang/multimodal_gen/test/unit/test_sdpa_cudnn_priority.py b/python/sglang/multimodal_gen/test/unit/test_sdpa_cudnn_priority.py new file mode 100644 index 000000000..5e7ebba1e --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_sdpa_cudnn_priority.py @@ -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()))