From 37ed10bd2446ea64c5145a1dcbaa326241ae12f5 Mon Sep 17 00:00:00 2001 From: Mick Date: Sun, 14 Jun 2026 14:55:13 +0800 Subject: [PATCH] [diffusion] UX: reduce attention backend log noise (#28169) --- .../runtime/layers/attention/selector.py | 66 ++++++++++++++----- .../multimodal_gen/runtime/platforms/cuda.py | 19 ------ 2 files changed, 51 insertions(+), 34 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/selector.py b/python/sglang/multimodal_gen/runtime/layers/attention/selector.py index 2f75b715e..07f61baf6 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/selector.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/selector.py @@ -67,6 +67,7 @@ forced_attn_backend: AttentionBackendEnum | None = None class ComponentAttnBackendContext(NamedTuple): backend: AttentionBackendEnum | None component_name: str | None + selected_backends: dict[str, str | None] component_attn_backend_context: ContextVar[ComponentAttnBackendContext | None] = ( @@ -111,6 +112,40 @@ def get_component_attn_backend_name() -> str | None: return context.component_name if context is not None else None +def _record_component_attn_backend(backend_name: str, reason: str | None) -> bool: + context = get_component_attn_backend_context() + if context is None or context.component_name is None: + return False + + existing_reason = context.selected_backends.get(backend_name) + if backend_name not in context.selected_backends or existing_reason is None: + context.selected_backends[backend_name] = reason + return True + + +def _log_component_attn_backend_summary( + context: ComponentAttnBackendContext | None, +) -> None: + if ( + context is None + or context.component_name is None + or not context.selected_backends + ): + return + + backend_parts = [] + for backend_name, reason in context.selected_backends.items(): + if reason: + backend_parts.append(f"{backend_name} ({reason})") + else: + backend_parts.append(backend_name) + + logger.info_once( + f"Attention backends for {context.component_name}: " + f"{', '.join(backend_parts)}" + ) + + def get_attn_backend( head_size: int, dtype: torch.dtype, @@ -141,23 +176,21 @@ def get_attn_backend( f"Available options are: {[e.name.lower() for e in AttentionBackendEnum]}" ) - component_name = get_component_attn_backend_name() - backend_not_specified = selected_backend is None + constraint_backend = None + if selected_backend is None and len(be_tuple) == 1: + constraint_backend = be_tuple[0].name.lower() + attention_backend_cls = _cached_get_attn_backend( head_size, dtype, be_tuple, selected_backend, ) - if component_name: - backend_name = attention_backend_cls.get_enum().name.lower() - if backend_not_specified: - logger.info_once( - f"Attention backend not specified for {component_name}, " - f"using {backend_name} backend for {component_name}" - ) - else: - logger.info_once(f"Using {backend_name} backend for {component_name}") + + backend_name = attention_backend_cls.get_enum().name.lower() + reason = "component constraint" if backend_name == constraint_backend else None + if not _record_component_attn_backend(backend_name, reason): + logger.info_once(f"Using {backend_name} attention backend") return attention_backend_cls @@ -178,9 +211,10 @@ def _cached_get_attn_backend( pass elif selected_backend is None and len(supported_attention_backends) == 1: selected_backend = next(iter(supported_attention_backends)) - elif selected_backend is None: - logger.debug("Attention backend not specified") - elif selected_backend not in supported_attention_backends: + elif ( + selected_backend is not None + and selected_backend not in supported_attention_backends + ): supported_attention_backends_str = [ supported_attention_backend.__str__() for supported_attention_backend in supported_attention_backends @@ -212,11 +246,13 @@ def component_attn_backend_context_manager( return token = component_attn_backend_context.set( - ComponentAttnBackendContext(attn_backend, component_name) + ComponentAttnBackendContext(attn_backend, component_name, {}) ) try: yield finally: + context = component_attn_backend_context.get() + _log_component_attn_backend_summary(context) component_attn_backend_context.reset(token) diff --git a/python/sglang/multimodal_gen/runtime/platforms/cuda.py b/python/sglang/multimodal_gen/runtime/platforms/cuda.py index 7a0b94310..d3d605170 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/cuda.py +++ b/python/sglang/multimodal_gen/runtime/platforms/cuda.py @@ -233,8 +233,6 @@ class CudaPlatformBase(Platform): SlidingTileAttentionBackend, ) - logger.info("Using Sliding Tile Attention backend") - return "sglang.multimodal_gen.runtime.layers.attention.backends.sliding_tile_attn.SlidingTileAttentionBackend" except ImportError as e: logger.error( @@ -251,8 +249,6 @@ class CudaPlatformBase(Platform): SageAttentionBackend, ) - logger.info("Using Sage Attention backend") - return "sglang.multimodal_gen.runtime.layers.attention.backends.sage_attn.SageAttentionBackend" except ImportError as e: logger.info(e) @@ -266,7 +262,6 @@ class CudaPlatformBase(Platform): SageAttention3Backend, ) - logger.info("Using Sage Attention 3 backend") return "sglang.multimodal_gen.runtime.layers.attention.backends.sage_attn3.SageAttention3Backend" except ImportError as e: logger.info(e) @@ -282,8 +277,6 @@ class CudaPlatformBase(Platform): VideoSparseAttentionBackend, ) - logger.info("Using Video Sparse Attention backend") - return "sglang.multimodal_gen.runtime.layers.attention.backends.video_sparse_attn.VideoSparseAttentionBackend" except ImportError as e: logger.error( @@ -309,7 +302,6 @@ class CudaPlatformBase(Platform): SparseVideoGen2AttentionBackend, ) - logger.info("Using Sparse Video Gen 2 (SAP) Attention backend") return "sglang.multimodal_gen.runtime.layers.attention.backends.sparse_video_gen_2_attn.SparseVideoGen2AttentionBackend" except ImportError as e: logger.error( @@ -329,8 +321,6 @@ class CudaPlatformBase(Platform): VMOBAAttentionBackend, ) - logger.info("Using Video MOBA Attention backend") - return "sglang.multimodal_gen.runtime.layers.attention.backends.vmoba.VMOBAAttentionBackend" except ImportError as e: logger.error( @@ -340,23 +330,18 @@ class CudaPlatformBase(Platform): "Video MoBA Attention backend is not installed. " ) from e elif selected_backend == AttentionBackendEnum.AITER: - logger.info("Using AITer backend") return "sglang.multimodal_gen.runtime.layers.attention.backends.aiter.AITerBackend" elif selected_backend == AttentionBackendEnum.TORCH_SDPA: - logger.info("Using Torch SDPA backend") return "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend" elif selected_backend == AttentionBackendEnum.SLA_ATTN: - logger.info("Using Sparse Linear Attention backend") return "sglang.multimodal_gen.runtime.layers.attention.backends.sparse_linear_attn.SparseLinearAttentionBackend" elif selected_backend == AttentionBackendEnum.SAGE_SLA_ATTN: - logger.info("Using Sage Sparse Linear Attention backend") return "sglang.multimodal_gen.runtime.layers.attention.backends.sparse_linear_attn.SageSparseLinearAttentionBackend" elif selected_backend == AttentionBackendEnum.FA2: from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn_2 import ( # noqa: F401 FlashAttention2Backend, ) - logger.info("Using FlashAttention2 backend") return "sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn_2.FlashAttention2Backend" elif selected_backend in [ AttentionBackendEnum.FA, @@ -438,12 +423,8 @@ class CudaPlatformBase(Platform): target_backend = AttentionBackendEnum.TORCH_SDPA if target_backend == AttentionBackendEnum.TORCH_SDPA: - logger.info("Using Torch SDPA backend") - return "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend" - logger.info("Using FlashAttention (FA3 for hopper, FA4 for blackwell) backend") - return "sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn.FlashAttentionBackend" @classmethod