执行摘要

• 一句话：统一混合 SWA 内存池改用共享字节预算并重构准入记账
• 推荐动作：值得精读。重点学习三处设计决策：一是 create_prefill_budget 工厂把准入记账下沉到分配器，消除调度器对池类型的 if/else 分支；二是 can_reserve/reclaim_plan 联合判定配合穷举式派生属性测试，是防御共享字节记账回归的有效范式；三是兄弟类重构解决了 Mamba 三池错误继承两池策略的抽象问题。遗留跟进项（rationale 注释、PrefillDelayer 水线、动态口径文档化）应在后续 PR 中闭环。

功能与动机

PR body 明确说明：统一的全注意力池与滑窗池共享同一字节 arena，但静态 per-pool 容量门会拒绝那些本可借用对方空闲字节的分配，容量报告必须跟随同样的动态分配状态。本 PR 是 #36403 合并后的 follow-up 栈第一项，目标是让 hybrid-SWA unified memory 的准入与回收都基于真实可用的共享字节空间，而不是两个互不感知的静态分区。

实现拆解

1. 新增 allocator-owned 准入预算抽象：新建 python/sglang/srt/mem_cache/prefill_budget.py，定义 PrefillBudget / SWAPrefillBudget / SharedSWAPrefillBudget 三级预算，承载 check_prefill、reserve、fit_chunk、available_chunk_tokens 等准入记账；同时在 allocator/base.py、allocator/swa.py、allocator/hisparse.py 增加 create_prefill_budget 工厂方法，让每种分配器返回与自己容量口径匹配的预算对象，调度器不再感知具体池类型。
2. 重构 unified hybrid-SWA 分配器：unified_hybrid_swa.py 拆出 UnifiedSWAAllocatorBase 基类（共享虚拟 ID/生命周期/索引翻译），UnifiedSWATokenToKVPoolAllocator 与 UnifiedMambaSWATokenToKVPoolAllocator 改为兄弟类，避免 Mamba 三池错误继承两池的联合字节预算判定；新增 can_reserve / reclaim_plan / _fits_page_demand / current_full_capacity / current_swa_capacity，evict_from_tree_cache（mem_cache/common.py）改为按 reclaim_plan 做联合回收。
3. 池尺寸显式字节预算：model_executor/pool_configurator.py 新增 _unified_pool_bytes 与 _make_pool_config，MemoryPoolConfig 增加 unified_memory_pool_bytes 字段，kv_cache_configurator.py 据此直接给 arena 定字节预算；draft worker 显式清空该字段，避免继承 target 预算；unified_memory_pool.py 工厂收敛为单一 total_bytes 参数，移除冗余的 unified_total_bytes 双输入与互斥校验。
4. 调度与统计配套：schedule_policy.py 的 PrefillAdder 改为持有 memory_budget 并委托全部预算计算，删除约 246 行手写 rem_*_token_offset 逻辑；scheduler.py 的 enqueue 裁剪拆出 _clip_unified_swa_max_new_tokens，统一“FULL 需求只取整一次、由 can_reserve 做页对齐”的规则；pool_stats_observer.py 改用动态容量口径报告；kv_index_translator.py 新增 capture_token_capacity 修复 capturer 索引契约。
5. 测试与清理：新增 test_prefill_memory_budget.py、test_unified_byte_budget_sizing.py 等 13+ 个测试文件，扩展 test_prefill_adder.py、test_pool_configurator.py、test_multi_ended_allocator.py，覆盖未对齐 final chunk、ignore-EOS 入空池、reclaim_plan 穷举、draft 预算隔离、派生属性边界；删除死代码 recover_swa_with_locked_full，并将 HiSparse fixture 适配 Python 3.10 的 with 语法。

关键文件：

• python/sglang/srt/mem_cache/prefill_budget.py（模块 准入预算；类别 source；类型 core-logic；符号 estimate_swa_kv_tokens, PrefillBudget, SWAPrefillBudget, SharedSWAPrefillBudget）: 新增的准入预算抽象核心文件：定义 PrefillBudget / SWAPrefillBudget / SharedSWAPrefillBudget 三级预算与 estimate_swa_kv_tokens 估算器，承载共享字节预算下的 check_prefill、reserve、fit_chunk 全部逻辑，是本 PR 架构转变的主入口。
• python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py（模块 共享池；类别 source；类型 core-logic；符号 UnifiedSWAAllocatorBase, UnifiedSWATokenToKVPoolAllocator, UnifiedMambaSWATokenToKVPoolAllocator, can_reserve）: 分配器重构核心：拆出 UnifiedSWAAllocatorBase，两池与 Mamba 三池变兄弟类；新增 can_reserve / reclaim_plan / current_*_capacity 联合判定与动态容量报告，是共享字节预算的落地载体。
• python/sglang/srt/managers/schedule_policy.py（模块 调度策略；类别 source；类型 core-logic；符号 PrefillAdder, memory_budget, rem_total_tokens, cur_rem_tokens）: PrefillAdder 从手写 rem_*_token_offset 改为持有 allocator 创建的 memory_budget，删除约 246 行分支记账逻辑，是准入路径的主要简化点。
• python/sglang/srt/model_executor/pool_configurator.py（模块 池配置；类别 source；类型 data-contract；符号 MemoryPoolConfig, _unified_pool_bytes, _make_pool_config, _max_unified_full_tokens）: 新增 unified_memory_pool_bytes 数据契约：_unified_pool_bytes 计算显式共享字节预算，_make_pool_config 统一构造 MemoryPoolConfig，draft 与 target 的预算隔离从这里开始。
• python/sglang/srt/mem_cache/allocator/base.py（模块 分配器；类别 source；类型 core-logic；符号 create_prefill_budget, max_new_tokens_for_memory, evict_to_free_tokens）: 定义 create_prefill_budget 工厂契约与 max_new_tokens_for_memory、evict_to_free_tokens 新语义，是所有分配器对准入预算抽象的接入点。
• test/registered/unit/mem_cache/test_prefill_memory_budget.py（模块 单元测试；类别 test；类型 test-coverage；符号 TestSharedPrefillMemoryBudget, TestSharedPrefillAdmission, test_pending_batch_cannot_spend_shared_bytes_twice, test_unaligned_final_chunk_makes_progress）: 新增 300 行核心回归：共享预算的 pending 请求不能重复花共享字节、混合 decode 两侧预留、前缀锁影响准入、未对齐 final chunk 与 ignore-EOS 入空池等，是 P1 修复的直接保障。
• test/registered/unit/managers/test_prefill_adder.py（模块 单元测试；类别 test；类型 test-coverage；符号 create_shared_adder, create_shared_req, test_shared_admission_reserves_all_pending_requests, test_shared_admission_rechecks_after_prefix_lock）: 扩展 172 行：覆盖共享准入预留全部 pending 请求、前缀锁后重新判定、decode 消耗 chunk 预算时 continuation 延迟，以及 memory_budget 偏移断言。
• python/sglang/srt/mem_cache/kv_cache_configurator.py（模块 缓存配置；类别 source；类型 core-logic；符号 hybrid_swa_token_capacity, _init_unified_swa_pools）: 承载 hybrid_swa_token_capacity 与 draft 预算清除逻辑：effective_max_total_num_tokens 的容量策略从 model_runner 迁移至此，draft worker 不再继承 target 的 unified_memory_pool_bytes。

关键符号：estimate_swa_kv_tokens, PrefillBudget, SWAPrefillBudget, SharedSWAPrefillBudget, create_prefill_budget, can_reserve, reclaim_plan, current_full_capacity, current_swa_capacity, _clip_unified_swa_max_new_tokens, hybrid_swa_token_capacity, _unified_pool_bytes, _make_pool_config, evict_to_free_tokens, max_new_tokens_for_memory

关键源码片段

python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py

分配器重构核心：拆出 UnifiedSWAAllocatorBase，两池与 Mamba 三池变兄弟类；新增 can_reserve / reclaim_plan / current_*_capacity 联合判定与动态容量报告，是共享字节预算的落地载体。

class UnifiedSWAAllocatorBase(SWATokenToKVPoolAllocator):
    """共享 FULL/SWA 虚拟 ID、分配生命周期与索引翻译的公共基类
​
    两个具体子类（两池 / Mamba 三池）各自定义容量策略，避免
    Mamba+SWA 池错误继承两池的联合字节预算判定。
    """
​
    # 动态容量口径：当前可用 + 已分配，随对侧实际字节占用浮动
    # 这取代了旧的静态 per-layer 上限作为统计分母
    @property
    def current_full_capacity(self) -> int:
        return self.full_available_size() + self.full_attn_allocator.allocated_count()
​
    @property
    def current_swa_capacity(self) -> int:
        return self.swa_available_size() + self.swa_attn_allocator.allocated_count()
​
    # 每侧容量视图：静态守恒上限约束借贷方，schedulable 视图约束
    # 向共享间隙增长的一侧，两者取小避免过承诺
    def full_available_size(self) -> int:
        return min(
            self._conserve_full_available_size(),
            self.schedulable_full_available_size(),
        )

python/sglang/srt/managers/schedule_policy.py

PrefillAdder 从手写 rem_*_token_offset 改为持有 allocator 创建的 memory_budget，删除约 246 行分支记账逻辑，是准入路径的主要简化点。

# PrefillAdder 不再自己维护 rem_total_token_offset / rem_swa_token_offset 两套
# 手写偏移，而是把准入记账整体委托给 allocator 创建的预算对象
self.memory_budget = token_to_kv_pool_allocator.create_prefill_budget(
    tree_cache, num_mixed_decode_tokens=num_mixed_decode_tokens
)
​
if running_batch is not None:
    # 运行中请求的估算 token 占用并入预算的 total 偏移
    self.memory_budget.total_offset += sum(
        self._get_running_request_total_token_offset(r)
        for r in running_batch.reqs
    )
​
@property
def rem_total_tokens(self):
    return self.memory_budget.remaining_total
​
@property
def cur_rem_tokens(self):
    return self.memory_budget.remaining_current
​
def _check_prefill_budget(self, req, *, extend_input_len, total_tokens, swa_host_hit_length):
    # 准入统一走预算对象的 check_prefill：fixed / shared / ring 三种
    # 分配器各自实现容量口径，调度器不再感知具体池类型
    return self.memory_budget.check_prefill(
        extend_input_len=extend_input_len,
        total_tokens=total_tokens,
        max_new_tokens=self._swa_new_tokens(req),
        input_tokens=len(req.full_untruncated_fill_ids),
        swa_host_hit_length=swa_host_hit_length,
        chunk_limit=self.rem_chunk_tokens,
    )

评论区精华

评审由 ch-wan 与 merrymercy 多轮深度跟进，核心交锋如下：

• ch-wan 的第一轮总结："核心记账是自洽的……残余风险与其说是'此方案不可行'，不如说是'新数学几乎没有被测试'。" 作者最初保持测试中立，后在 c41b06843c 补齐派生属性与空池字节边界覆盖。
• merrymercy 提问："do we have a e2e unified hybrid swa memory pool test with an engine launch?" 作者确认 test_gemma4_unified_swa_virtual_ids.py 会启动真实 gemma-4-E2B-it 引擎并跑 200 条 GSM8K，但也明确 "I have not rerun this engine-launch test for the newly pushed fixes"，随后通过 /rerun-test 补跑一次通过。
• merrymercy 指出："class UnifiedMambaSWATokenToKVPoolAllocator(UnifiedSWATokenToKVPoolAllocator) is wrong abstraction"，作者在 f20860697232 中把两个分配器改为 UnifiedSWAAllocatorBase 下的兄弟类。
• ch-wan 报告 P2：enqueue 裁剪先对 prompt 取整、can_reserve 再取整一次，导致 page_size=4 时 FULL 需求按 13 而非 10 记账，可行生成被裁剪到 0 token；作者修复为单一取整并加 test_unified_budget_rounds_prompt_and_decode_together。
• merrymercy 报告两个 P1：未对齐 final chunk 与 ignore-EOS 入空池在共享预算下被永久拒绝（fit_chunk 多计一页 FULL），作者在 1f3b496c23 修复并新增 TestSharedPrefillAdmission 四个真实 Req/PrefillAdder 边界回归。

• 未闭环项：ch-wan 要求把 min(extend+decode, window) 与 chunk 收缩的 rationale 注释移入 estimate_swa_kv_tokens，作者承认仍待补；统计分母变为动态容量后，PrefillDelayer 水线是否需要重调也未处理，作者标注为后续跟进。

• 双重取整导致可行生成被裁剪为 0 token (correctness): 修复为 _clip_unified_swa_max_new_tokens 单一取整、can_reserve 统一舍入，并加 test_unified_budget_rounds_prompt_and_decode_together 回归。

• 未对齐 final chunk 与 ignore-EOS 入空池被永久拒绝 (correctness): 1f3b496c23 修复：共享预算下统一传未取整长度，新增 TestSharedPrefillAdmission 四个真实 Req/PrefillAdder 边界回归，page_size 覆盖 4/64。
• estimate_swa_kv_tokens 的 allocation_limit=0 语义重载 (design): 先加 include_allocated_tail 标志，最终在 ab2ea29 删除该标志与 _swa_reserved_tokens，统一为 allocation_limit 单一参数。
• reclaim_plan 对 SWA>FULL 需求直接返回 None 导致队头阻塞 (correctness): 4d1fa0ff 移除 SWA>FULL 的硬拒绝，按供需处理；test_restore_swa_without_allocating_more_full 与穷举式 reclaim 测试覆盖。
• draft worker 继承 target 的共享字节预算 (correctness): ab2ea29 在 draft config 显式清除 unified_memory_pool_bytes，并加 test_draft_does_not_inherit_target_shared_byte_budget 回归。
• UnifiedMamba 池继承两池分配器的抽象错误 (design): f208606 拆出 UnifiedSWAAllocatorBase 作为公共基类，两池与三池变兄弟类，原 supports_asymmetric_reservation 标志机制被移除。
• 派生属性测试缺口：_max_unified_full_tokens、can_reserve、reclaim_plan (testing): c41b068 补 derived-property 覆盖（最大 page-aligned full 容量、空池 packed-byte 边界）；test_multi_ended_allocator 增加 reclaim_plan 穷举对照。
• 动态容量口径对 token usage 指标与 capturer 的影响 (question): 动态口径保留，capturer 索引修复；PrefillDelayer 政策未改，水线重调作为运维跟进项。
• estimate_swa_kv_tokens 丢失 rationale 注释 (documentation): 作者同意并确认仍未恢复，作为文档跟进项开放。

风险与影响

• 风险：技术风险集中在以下几处：
  1. 取整语义回归：本 PR 历史上有两个 P1/P2 缺陷都源于“FULL 需求双重取整”。修复后统一为 can_reserve 单一舍入口径，但该契约依赖调用方传入未取整长度；后续新增调用者（如 HiCache/PD 路径在 #37496 才可达）若沿用旧的 paged_input 传参方式，会重新触发“enqueue 接受但 PrefillAdder 永远拒绝”的队头阻塞。
  2. 动态容量口径对指标与缓冲的影响：pool_stats_observer 的分母从静态 per-layer 上限变为 available + allocated，同一负载 token usage 从 0.38 变为 0.21；PrefillDelayer 的 waterwall 与 capturer 尺寸都可能需要重调，capturer 索引契约已由 kv_index_translator.capture_token_capacity 修复，但 GPU 环境下的长稳表现未充分验证。
  3. 启动期行为收紧：server_args._handle_unified_memory_pool 曾把 --enable-unified-memory 与 PP/mixed-chunk 组合一律断言禁止，后放宽为仅 PD 场景限制 PP；切换过程中若用户配置恰好踩线，会在启动阶段直接失败。
  4. CI 未全绿即合入：PR body 中三个 CI 看板均为小红叉；合入主要依赖本地 331 passed/722 subtests 与单次 gemma-4 E2E，未覆盖 multi-node PD、HiCache 组合，属于覆盖盲区。
  5. 冻结文件规则越界：model_runner.py 的 effective_max_total_num_tokens 一度直接内联 unified 容量策略，后移至 kv_cache_configurator.hybrid_swa_token_capacity()；新契约暴露后，后续调用方仍需遵守 frozen-core 边界。
     - 影响：影响面限定在开启 --enable-unified-memory 的 hybrid-SWA 部署：调度准入、预取回收、统计口径、capturer 容量、draft worker 池尺寸均随共享字节预算改变；UMP-off 的静态分区路径保持上游行为。对运维而言，token usage 类指标含义变化，--prefill-delayer-token-usage-low-watermark 等参数可能需要重调。对团队而言，本 PR 确立了 allocator-owned budget 的契约、派生属性测试风格与兄弟类继承模式，为后续 UMP 堆栈（#36730/#36731/#37496）打下基础。
     - 风险标记：双重取整曾致 P1/P2 缺陷, 动态容量口径需重调水线, CI 未全绿即合入, 冻结文件规则越界后修复, 文档跟进项未闭环

关联脉络

• PR #36403 Unified memory pool groundwork（前序，已合并）: PR body 明确说明本 PR 构建在其之上（"Builds on #36403, which is now merged"），是 unified memory 字节预算线的直接前序。
• PR #36730 Follow-up stack PR（讨论中提及）: review 评论确认本 PR 的修复已通过 #36730 传播，属于同一 UMP 共享预算演进栈。
• PR #36731 Follow-up stack PR（讨论中提及）: review 评论确认本 PR 的修复已通过 #36731 传播，属于同一 UMP 共享预算演进栈。
• PR #37496 HiCache allow-list 相关 PR（讨论中提及）: ch-wan 指出 #37496 的 HiCache allow-list 使 reclaim_plan 的 SWA>FULL 缺陷变得可达，说明该 PR 与本 PR 的联合回收逻辑直接交互。