minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] UX: reduce attention backend log noise (#28169)

Browse Source
This commit is contained in:
Mick
2026-06-14 14:55:13 +08:00
committed by GitHub
parent cc72e2bd8c
commit 37ed10bd24
2 changed files with 51 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/layers/attention/selector.py
+51 -15
View File
@@ -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)
python/sglang/multimodal_gen/runtime/platforms/cuda.py
-19
View File
@@ -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