执行摘要

• 一句话：支持 diffusion 组件级 attention backend 覆盖
• 推荐动作：值得阅读，特别是 ContextVar 为基础的组件化上下文注入模式，以及命名解析与回退策略。若你负责扩散推理优化，此 PR 提供了灵活的扩展基准。

功能与动机

在 diffusion 管线中，不同组件（如 text encoder、transformer）可能适合不同的 attention backend（例如 torch_sdpa 更适合 CPU、flash attention 更适合 GPU）。之前只能全局设置 --attention-backend，无法细粒度控制。该 PR 允许用户通过 --component-attention-backends 按组件指定，灵活调整性能与精度。

实现拆解

1. CLI 参数与解析：在 ServerArgs 添加 component_attention_backends 字段，实现 _normalize_attention_backend_name、_parse_component_attention_backend_map、_normalize_component_attention_backends 和 resolve_component_attention_backend 等方法，支持字典、逗号分隔字符串及点号 CLI 参数。
2. 上下文注入：在 selector.py 中定义 ComponentAttnBackendContext NamedTuple 和 ContextVar，提供 get_component_forced_attn_backend 等 getter。get_attn_backend 新增 selected_attention_backend 参数，优先级调整。
3. 组件后端推断：denoising.py 的 _infer_transformer_attention_backend 方法扫描 transformer 子模块的 backend 属性，自动推测当前组件使用的 backend，替代原先的 TODO hack。
4. 加载时注入：在 composed_pipeline_base.py 和 component_loader.py 中，加载组件前调用 resolve_component_attention_backend 并通过 component_attn_backend_context_manager 注入上下文。
5. 测试与文档：新增 4 个单元测试检查解析和错误处理；更新两处文档说明 CLI 用法和设计思路。

关键文件：

• python/sglang/multimodal_gen/runtime/server_args.py（模块 参数解析；类别 source；类型 core-logic；符号 _normalize_attention_backend_name, _parse_component_attention_backend_map, _normalize_component_attention_backends, resolve_component_attention_backend）: 新增 component_attention_backends 字段及解析、规范化、查找方法，是功能入口
• python/sglang/multimodal_gen/runtime/layers/attention/selector.py（模块 注意力选择器；类别 source；类型 core-logic；符号 ComponentAttnBackendContext, get_component_attn_backend_context, get_component_forced_attn_backend, get_component_attn_backend_name）: 引入 ContextVar 组件上下文和 selected_attention_backend 参数，实现动态覆盖
• python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py（模块 去噪阶段；类别 source；类型 core-logic；符号 _infer_transformer_attention_backend）: 添加 _infer_transformer_attention_backend 方法推断组件后端，替代 TODO hack
• python/sglang/multimodal_gen/test/unit/test_server_args.py（模块 测试；类别 test；类型 test-coverage；符号 test_component_attention_backends_are_normalized, test_component_attention_backend_lookup, test_invalid_component_attention_backend_raises, test_dynamic_component_attention_backend_cli_args）: 新增 4 个单元测试覆盖 CLI 解析、规范化、错误输入和动态参数
• python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py（模块 流水线基础；类别 source；类型 dependency-wiring）: 在 load_modules 中注入组件 attention backend context
• python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py（模块 组件加载器；类别 source；类型 dependency-wiring）: 在 load_customized/load_native 中包装上下文管理器
• docs_new/docs/sglang-diffusion/api/cli.mdx（模块 文档；类别 other；类型 entrypoint）: 更新 CLI 文档，说明 --component-attention-backends 用法
• docs_new/docs/sglang-diffusion/attention_backends.mdx（模块 文档；类别 other；类型 core-logic）: 更新 attention backends 文档，说明组件级覆盖
• docs/diffusion/api/cli.md（模块 文档；类别 docs；类型 documentation）: 更新旧版 CLI 文档
• docs/diffusion/performance/attention_backends.md（模块 文档；类别 docs；类型 documentation）: 更新旧版性能文档

关键符号：_normalize_attention_backend_name, _parse_component_attention_backend_map, _normalize_component_attention_backends, resolve_component_attention_backend, _extract_component_attention_backends, ComponentAttnBackendContext, get_component_attn_backend_context, get_component_forced_attn_backend, get_component_attn_backend_name, _cached_get_attn_backend, _infer_transformer_attention_backend, component_attn_backend_context_manager

关键源码片段

python/sglang/multimodal_gen/runtime/server_args.py

新增 component_attention_backends 字段及解析、规范化、查找方法，是功能入口

# server_args.py - 标准化与解析
@staticmethod
def _normalize_attention_backend_name(backend: str) -> str:
    # 标准化 backend 名称，fa3/fa4 映射为 fa
    if not isinstance(backend, str):
        raise ValueError('Attention backend name must be a string')
    normalized = backend.strip().lower()
    if normalized in ('fa3', 'fa4'):
        normalized = 'fa'
    try:
        return AttentionBackendEnum[normalized.upper()].name.lower()
    except KeyError:
        raise ValueError(
            f'Invalid attention backend {backend}. '
            f'Available options: {[e.name.lower() for e in AttentionBackendEnum]}'
        ) from None
​
@classmethod
def _normalize_component_attention_backends(
    cls, value: dict[str, str] | str | None
) -> dict[str, str]:
    # 先解析为字典，然后将每个 backend 名称规范化
    raw = cls._parse_component_attention_backend_map(value)
    normalized: dict[str, str] = {}
    for component, backend in raw.items():
        # 将组件名中的 '-' 替换为 '_'，以便与 Python 标识符对齐
        component_name = component.strip().replace('-', '_')
        if not component_name:
            raise ValueError('Component attention backend key must not be empty')
        normalized[component_name] = cls._normalize_attention_backend_name(backend)
    return normalized
​
def resolve_component_attention_backend(
    self, *component_names: str | None
) -> tuple[AttentionBackendEnum | None, str | None]:
    # 按传入的组件名称顺序查找覆盖，返回 (enum, 匹配键 ) 或 (None, None)
    for name in component_names:
        if name is None:
            continue
        backend_str = self.component_attention_backends.get(name)
        if backend_str is not None:
            try:
                backend = AttentionBackendEnum[backend_str.upper()]
                return backend, name
            except KeyError:
                raise ValueError(
                    f'Invalid attention backend {backend_str} for component {name}.'
                )
    return None, None

python/sglang/multimodal_gen/runtime/layers/attention/selector.py

引入 ContextVar 组件上下文和 selected_attention_backend 参数，实现动态覆盖

# selector.py - 组件上下文定义与优先级逻辑
class ComponentAttnBackendContext(NamedTuple):
    backend: AttentionBackendEnum | None
    component_name: str | None
​
# 使用 ContextVar 存储当前线程的组件上下文，默认 None
component_attn_backend_context: ContextVar[ComponentAttnBackendContext | None] = (
    ContextVar('component_attn_backend_context', default=None)
)
​
def get_component_forced_attn_backend() -> AttentionBackendEnum | None:
    # 获取当前组件上下文强制使用的 backend，若无则返回 None
    context = component_attn_backend_context.get()
    return context.backend if context is not None else None
​
def get_component_attn_backend_name() -> str | None:
    # 获取当前组件的名称，用于日志
    context = component_attn_backend_context.get()
    return context.component_name if context is not None else None
​
def get_attn_backend(
    head_size: int,
    dtype: torch.dtype,
    supported_attention_backends: set[AttentionBackendEnum] | None = None,
    selected_attention_backend: AttentionBackendEnum | None = None,
) -> type[AttentionBackend]:
    # 构建缓存键（略）
    # 优先级链：显式传入 > 全局强制 > 组件上下文 > server_args
    selected_backend = selected_attention_backend or get_global_forced_attn_backend()
    if selected_backend is None:
        selected_backend = get_component_forced_attn_backend()
    if selected_backend is None:
        server_args = get_global_server_args()
        if server_args.attention_backend is not None:
            try:
                selected_backend = AttentionBackendEnum[
                    server_args.attention_backend.upper()
                ]
            except KeyError:
                raise ValueError(
                    f'Invalid attention backend {server_args.attention_backend}'
                )
    # 继续设备自动选择（略）

评论区精华

本次 PR 无实质人工讨论。仅包含 gemini-code-assist 的自动 review，无反馈需要处理。

• 暂无高价值评论线程

风险与影响

• 风险：
  1. ContextVar 在多线程或异步任务中可能泄漏上下文，需要确保上下文管理器正确进出。
  2. 组件名称规范化（- 替换为 _）可能引起用户混淆，且 fa3/fa4 映射为 fa 会丢失原始信息，未来需保持映射更新。
  3. resolve_component_attention_backend 采用精确匹配，组件名称拼写差异可能导致覆盖失效。
  4. 目前仅在 diffusion 模块生效，对核心调度器无影响。
  5. 测试覆盖基本路径，但未覆盖异常嵌套、多组件并行加载等场景。
     - 影响：对用户：扩散模块用户可以精细控制每个组件的 attention backend，有利于异构硬件或精度要求下的性能调优。对系统：新增的 ContextVar 机制不会影响现有全局 backend 选择，向后兼容。对团队：需要维护新增 API 和文档，并且未来扩展其他组件时需要保持一致性。
     - 风险标记：ContextVar 上下文泄漏风险, 组件名称规范化歧义, fa3/fa4 归一化丢失信息, 精确匹配组件名, 缺少并行加载测试

关联脉络

• 暂无明显关联 PR