Prhub

#35932 Make draft attention backends extensible

原始 PR 作者 merrymercy 合并时间 2026-08-22 14:54 文件变更 2 提交数 2 评论 1 代码增减 +18 / -15

执行摘要

草稿注意力后端白名单迁移到共享注册表并开放扩展

PR body 指出:Out-of-tree integrations can register general attention backends, but draft workers keep a separate private allowlist. As a result, a draft-compatible external backend cannot opt into the common draft worker path. 因此需要把这个私有白名单开放为可扩展的共享注册表,让外部后端能在不修改核心逻辑的前提下复用草稿 worker 的公共路径。

适合快速阅读,作为理解 SGLang server_args 注册表扩展模式的小样例。若你正在维护外部注意力后端集成,建议关注 add_draft_attention_backend_choices 的调用时机(必须在 ServerArgs 构造前)。若你关心配置系统重构,本 PR 与近期大量 config 相关 PR 展示的注册表模式一致。

讨论亮点

该 PR 没有实质 review 讨论线程;唯一评论是作者触发的 /tag-and-rerun-ci。作者在 body 中说明 Accuracy/Speed Tests 不适用,因为内置后端的计算路径未改变。

实现拆解

  1. 在 python/sglang/srt/server_args.py 中新增模块级常量 DRAFT_ATTENTION_BACKEND_CHOICES,内容与原本的 _SUPPORTED_DRAFT_BACKENDS 完全一致(flashinfer/fa3/fa4/triton/ascend/trtllm_mha),并补充注释说明 trtllm_mha 仅适用于 decode-only 密集 MQA 草稿。
  2. 新增注册辅助函数 add_draft_attention_backend_choices(choices),对 DRAFT_ATTENTION_BACKEND_CHOICES 执行 extend,与已有的 add_attention_backend_choices / add_deterministic_attention_backend_choices 等扩展入口保持同一模式。
  3. 在 python/sglang/srt/speculative/draft_worker_common.py 中删除私有元组 _SUPPORTED_DRAFT_BACKENDS,改为从 server_args 导入 DRAFT_ATTENTION_BACKEND_CHOICES;_resolve_draft_attention_backend_fallback 中的白名单判断和告警日志均改用它。
  4. 测试配套:未新增单元测试文件,作者通过 pre-commit、py_compile 和 registry-to-resolver smoke test 验证注册与解析链路;该改动要求注册发生在 ServerArgs 构造/解析之前。
文件 模块 状态 重要度
python/sglang/srt/server_args.py 参数注册 modified 6.27
python/sglang/srt/speculative/draft_worker_common.py 投机解码 modified 5.87

关键符号

add_draft_attention_backend_choices _resolve_draft_attention_backend_fallback

关键源码片段

python/sglang/srt/server_args.py core-logic

核心变更文件:新增 DRAFT_ATTENTION_BACKEND_CHOICES 白名单和 add_draft_attention_backend_choices 扩展入口,使草稿注意力后端可以像普通注意力后端一样被外部注册。

# 草稿解码(draft decoding)可用的注意力后端白名单。
# 该列表与 ATTENTION_BACKEND_CHOICES 分离,因为草稿 worker 的
# 后端支持范围更小(例如 trtllm_mha 只支持 decode-only 的密集 MQA 草稿)。
# 外部集成可通过 add_draft_attention_backend_choices() 在 ServerArgs
# 构造 / 解析前扩展该列表,使自研后端能进入公共草稿 worker 路径。
DRAFT_ATTENTION_BACKEND_CHOICES = [
    'flashinfer',
    'fa3',
    'fa4',
    'triton',
    'ascend',
    'trtllm_mha',
]# 为草稿 worker 注册额外的注意力后端名称。
def add_draft_attention_backend_choices(choices):
    DRAFT_ATTENTION_BACKEND_CHOICES.extend(choices)
python/sglang/srt/speculative/draft_worker_common.py dependency-wiring

将草稿 worker 的私有白名单改为引用共享注册表,fallback 逻辑随之使用可扩展列表,是本 PR 的消费侧变更。

def _resolve_draft_attention_backend_fallback(
    *, server_args: ServerArgs, algo_label: str
) -> str:
    # 优先使用显式指定的草稿后端;未指定时从目标模型的后端解析结果中继承。
    draft_backend = server_args.speculative_draft_attention_backend
    if draft_backend is None:
        draft_backend, _ = server_args.get_attention_backends()
    if draft_backend is None:
        # 平台默认:AMD 用 triton,否则用 flashinfer
        return 'triton' if torch.version.hip else 'flashinfer'
​
    # 不在共享注册表中或未注册的后端,回退到平台默认值并告警。
    # 通过 DRAFT_ATTENTION_BACKEND_CHOICES 判断,而不是私有常量,
    # 这样外部注册的后端也能被识别,而不是被误判为不支持。
    if draft_backend not in DRAFT_ATTENTION_BACKEND_CHOICES:
        fallback = 'triton' if torch.version.hip else 'flashinfer'
        logger.warning(
            '%s draft worker only supports attention_backend in %s for now, '
            'but got %r. Falling back to %s.',
            algo_label,
            DRAFT_ATTENTION_BACKEND_CHOICES,
            draft_backend,
            fallback,
        )
        return fallback
    return draft_backend

评论区精华

没有提炼出高价值讨论线程

当前评论区没有形成足够清晰的争议点或结论,后续有更多讨论时会体现在这里。

风险与影响

主要风险集中在缺少正式回归测试;DRAFT_ATTENTION_BACKEND_CHOICES 是模块级可变列表,若外部代码在 ServerArgs 完成解析后再调用 add_draft_attention_backend_choices,新增项不会生效且无报错。此外,共享列表可被外部意外修改,不过这与 SGLang 现有注册函数风格一致。fallback 逻辑与改动前等价,因此内置后端行为无回归风险。

对内置用户无影响:后端解析顺序、默认值、回退逻辑均不变。对外部集成者有实际收益:自研的兼容草稿注意力后端现可通过注册函数进入公共草稿 worker 路径,无需 fork 或 patch SGLang。对团队而言,消除了 draft_worker_common.py 中重复定义的白名单,降低未来新增内置后端时漏改的风险。

缺少单元测试 注册时机敏感,解析后注册无效 低风险小改动

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论