Prhub

#33334 config: stop writing config onto the published ServerArgs at three sites

原始 PR 作者 ch-wan 合并时间 2026-08-03 12:22 文件变更 10 提交数 1 评论 1 代码增减 +136 / -51

执行摘要

三处配置不再写入已发布 ServerArgs,改为入参或 override

这是五部分系列的第一部分,目标是消除对已发布 ServerArgs 实例的进程级配置写入。PR body 明确指出:"Three post-resolution ServerArgs.override call-sites wrote config onto the published instance. None of them needed the instance: one write was redundant, and two carry a value the resolved-config readback reports"。这些写入是在配置已发布(只读)之后发生的,导致 resolved-config 读回与 seed 不一致:例如 XGrammar 回退后读回会声称 XGrammar 可用而 structured output 实际不可用。此前 #33238 以 stack 方式提交导致 GitHub 无法合并,本 PR 以单分支重开。

值得精读,特别是 flashinfer_gdn_prefill_default 纯函数化 + initialize_linear_attn_config 参数注入的写法,以及用 ratchet 测试基线来度量进程级配置写入点数量的做法;也是理解 sglang 配置分层(seed/override/readback)架构的入门样例。建议关注系列后续四个 PR 如何继续清除其余写入点。

讨论亮点

本 PR 的 review 讨论都在前置 PR #33238 上进行(PR body 说明代码与 #33238 最终版一致),本 PR 自身没有独立 review 评论。核心讨论点是:三个写入点都不需要 ServerArgs 实例本身——一个写入是冗余死代码,两个写入的值由 resolved-config 读回覆盖,因此迁移到 get_context().override 或纯函数返回值;测试桩从 SimpleNamespace + MagicMock 模拟 override 改为直接断言返回值,降低了桩与真实实现的分叉风险。

实现拆解

  1. SM100 GDN prefill 默认值改为纯函数返回值:在 python/sglang/srt/layers/attention/linear/gdn_backend.py 中,将 maybe_set_default_flashinfer_gdn_prefill(写入 ServerArgs 的 override)重构为 flashinfer_gdn_prefill_default,返回 "flashinfer" 或 None;所有早退分支改为 return None。在 python/sglang/srt/layers/attention/attention_registry.py 的 attn_backend_wrapper 中先调用该函数取得 prefill_default,非 None 时通过 get_context().override 记录(读回仍可见),再传入 initialize_linear_attn_config。python/sglang/srt/layers/attention/linear/utils.py 的 initialize_linear_attn_config 新增 prefill_default 参数,优先级为显式 flag > prefill_default > base backend,与原有“显式 --linear-attn-prefill-backend 优先”语义完全一致,分支顺序、设备/CUDA/dtype 条件和日志行均不变。
  2. XGrammar tokenizer 回退写入 context override:在 python/sglang/srt/constrained/base_grammar_backend.py 的 create_grammar_backend 中,将 server_args.override("grammar.import_fallback", grammar_backend="none") 改为 get_context().override(...);调用方仍通过 None 返回值感知回退,警告日志保留,但 resolved-config 读回现在真实反映 fallback。
  3. 删除 HiCache direct-IO 布局修正:python/sglang/srt/mem_cache/unified_radix_cache.py 的 init_hicache 删除了将 page_first 改写为 page_first_direct 的 fixup 代码块。该规范化在 ServerArgs.post_init 中已执行,且 init_hicache 只在 hierarchical cache 启用时运行——正是 _handle_hicache 重写 page_first 的条件,因此原 fixup 是死代码。
  4. 测试配套:新增 test/registered/unit/layers/attention/test_linear_attn_config.py 锁定时序优先级(显式 flag、SM100 默认、base backend、decode 不受影响、override 读回);test_gdn_prefill_backend_policy.py 改为断言纯函数返回值并移除 SimpleNamespace 桩上的 override MagicMock;test_base_grammar_backend.py 改为断言 get_exec().kernel.grammar_backend 与 resolved_server_args_dict;test_unified_radix_cache_unittest.py 三处 fixture 显式声明 hicache_mem_layout="page_first_direct"(因为 dummy model_path 跳过 post_init 解析);test_server_args_writer_ratchet.py 基线 34 → 31。
文件 模块 状态 重要度
python/sglang/srt/layers/attention/linear/gdn_backend.py GDN 后端 modified 7.06
python/sglang/srt/layers/attention/attention_registry.py 注意力注册表 modified 6.16
python/sglang/srt/layers/attention/linear/utils.py 线性注意力配置 modified 5.98
python/sglang/srt/mem_cache/unified_radix_cache.py 缓存器 modified 5.37
python/sglang/srt/constrained/base_grammar_backend.py 语法后端 modified 5.07
test/registered/unit/layers/attention/test_linear_attn_config.py 配置测试 added 7.42
test/registered/unit/layers/attention/test_gdn_prefill_backend_policy.py 策略测试 modified 5.53
test/registered/unit/constrained/test_base_grammar_backend.py 语法测试 modified 4.95
test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py 缓存测试 modified 3.92
test/registered/unit/test_server_args_writer_ratchet.py 配置回归 modified 3.82

关键符号

flashinfer_gdn_prefill_default initialize_linear_attn_config attn_backend_wrapper init_hicache create_grammar_backend

关键源码片段

python/sglang/srt/layers/attention/linear/gdn_backend.py core-logic

SM100 GDN prefill 默认值的核心重构:maybe_set_default_flashinfer_gdn_prefill 改为纯函数 flashinfer_gdn_prefill_default,不再写 ServerArgs 而是返回 backend 或 None

# python/sglang/srt/layers/attention/linear/gdn_backend.py
def flashinfer_gdn_prefill_default(model_runner: ModelRunner) -> Optional[str]:
    """在已验证的窄 SM100 GDN prefill 域内返回 FlashInfer,否则返回 None。"""
    args = model_runner.server_args
    if (
        args.linear_attn_prefill_backend is not None # 显式 flag 优先,不产生默认值
        or args.linear_attn_backend != "triton"
        or args.enable_page_major_kv_layout
        or not is_cuda()
        or torch.cuda.get_device_capability()[0] != 10
    ):
        return None
​
    cuda_version = torch.version.cuda
    chunk_size = args.chunked_prefill_size
    config = hybrid_gdn_config(model_runner.model_config)
    if (
        cuda_version is None
        or int(cuda_version.split(".", 1)[0]) < 13
        or args.enable_dynamic_chunking
        or chunk_size is None
        or not 1 <= chunk_size <= 8192
        or getattr(config, "linear_key_head_dim", None) != 128
        or getattr(config, "linear_value_head_dim", None) != 128
        or model_runner.req_to_token_pool.mamba_pool.mamba_cache.temporal.dtype
        != torch.bfloat16
    ):
        return None
​
    from sglang.srt.layers.attention.linear.kernels.gdn_flashinfer import (
        is_flashinfer_gdn_prefill_available,
    )
​
    if not is_flashinfer_gdn_prefill_available():
        return None
​
    # 纯函数只负责决策;副作用(override 记录)由调用方在 attention_registry 执行
    rank0_log("Defaulting SM100 GDN prefill backend to FlashInfer.")
    return "flashinfer"
python/sglang/srt/layers/attention/attention_registry.py dependency-wiring

将默认值函数接入 attn_backend_wrapper:通过 get_context().override 记录默认值并传入 initialize_linear_attn_config,是迁移后的副作用落点

# python/sglang/srt/layers/attention/attention_registry.py(attn_backend_wrapper 内)
check_environments()
prefill_default = None
if hybrid_gdn_config(runner.model_config) is not None and not is_npu():
    # 纯函数判定是否启用 SM100 GDN flashinfer 默认值
    prefill_default = flashinfer_gdn_prefill_default(runner)
if prefill_default is not None:
    # 默认值通过 context override 记录:resolved-config 读回可见,但 seed 保持启动参数
    get_context().override(
        "gdn_backend.sm100_flashinfer_default",
        linear_attn_prefill_backend=prefill_default,
    )
# 默认值作为参数传入,而不是写入 ServerArgs 后再读回
initialize_linear_attn_config(runner.server_args, prefill_default)
python/sglang/srt/layers/attention/linear/utils.py core-logic

initialize_linear_attn_config 新增 prefill_default 参数,优先级为显式 flag > 默认值 > base backend,是行为保持的关键

# python/sglang/srt/layers/attention/linear/utils.py
def initialize_linear_attn_config(
    server_args: ServerArgs, prefill_default: Optional[str] = None
):
    global LINEAR_ATTN_DECODE_BACKEND
    global LINEAR_ATTN_PREFILL_BACKEND
​
    base = server_args.linear_attn_backend
    decode = server_args.linear_attn_decode_backend or base
    # 优先级:显式 flag > SM100 默认值 > base backend,与旧行为等价
    prefill = server_args.linear_attn_prefill_backend or prefill_default or base
​
    LINEAR_ATTN_DECODE_BACKEND = LinearAttnKernelBackend(decode)
    LINEAR_ATTN_PREFILL_BACKEND = LinearAttnKernelBackend(prefill)
​
    rank0_log(f"Linear attention kernel backend: decode={decode}, prefill={prefill}")

评论区精华

前置 PR #33238 的 review 讨论继承 设计

本 PR 的代码与已关闭的 #33238 最终版完全一致,PR body 说明 review 讨论和每轮评论的 triage 都在 #33238 上进行。GitHub 将此前 chained-base 系列视为 stack,阻碍了 base retargeting 和除 async endpoint 外的所有合并路径。

结论:以单分支重开并合入 main,五部分系列按顺序后续合并;讨论结论继承自 #33238。 · 已解决

测试桩从 SimpleNamespace + MagicMock override 改为纯函数返回值断言 测试

test_gdn_prefill_backend_policy.py 原先在 make_runner 的 SimpleNamespace 桩上模拟 server_args.override 副作用,断言 override 被调用一次;重构后策略函数变为纯函数,测试直接断言返回值,移除了桩上的 MagicMock override 模拟。

结论:接受纯函数化:测试不再需要模拟写入副作用,降低桩与真实实现的分叉风险;test_linear_attn_config.py 新测试用真实 ServerArgs + context override 锁定优先级。 · 已解决

风险与影响

风险集中在 SM100 GDN prefill 默认值链路:flashinfer_gdn_prefill_default 的模块级 wiring(attention_registry 中的调用)只在 GPU/Blackwell 上执行,本地无法覆盖(PR body 自述);若 flashinfer_gdn_prefill_default 返回 None 的某个分支条件与旧的 maybe_set_default_flashinfer_gdn_prefill 不一致(例如 is_cuda() 或 device capability 判断的时序),会导致 SM100 上 prefill backend 从 flashinfer 静默回退到 triton,影响性能而非正确性。其次,HiCache direct-IO 布局 fixup 的删除依赖 post_init 一定执行的假设;测试 fixture 用 dummy model_path 跳过了 post_init,因此改为显式声明 page_first_direct——如果生产路径上存在绕过 post_init 的构造方式(如直接 setattr),布局可能不正确,但这种情况生产未见。XGrammar fallback 改动风险最低,只是写入目标从实例改为 context。

影响面为配置系统内部一致性:resolved_server_args_dict 读回不再与 seed 冲突,get_internal_state 报告真实生效的配置(grammar_backend、linear_attn_prefill_backend);writer ratchet 从 34 降至 31 表明移除了 3 个进程级写入点。对用户无行为变化(显式 flag 仍然优先、日志不变),对 SM100 GDN 模型是纯内部重构;对团队的意义是推进 ServerArgs 只读化、配置分层(seed/override/readback)的架构方向,是该五部分系列的第一块基石,后续四个 PR 将在此基础上继续清除写入点。

SM100 GDN 路径无 GPU 覆盖 依赖 __post_init__ 规范化假设 默认值静默回退影响性能

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论