执行摘要

• 一句话：修复混合SWA分块预填充的内存溢出问题，确保预留页面空间并正确处理资源不足时的延迟处理。
• 推荐动作：该PR值得精读，特别是add_chunked_req方法中的设计决策：如何平衡资源分配与OOM风险，以及延迟处理与回退策略的选择。关注预留页面空间的实现细节和测试用例的设计，这些展示了在资源受限场景下的稳健性处理模式。

功能与动机

根据PR body描述，动机是解决混合SWA路径中分块预填充的内存溢出问题。具体来说，PrefillAdder.add_chunked_req在混合SWA路径中计算剩余令牌时，未考虑alloc_extend需要为每个请求预留extend_num_tokens + page_size的空间，导致当rem_swa_tokens == extend_num_tokens时，调度器分配全部SWA预算，但alloc_extend无法获取额外页面而OOM。此外，当_rem_tokens <= 0时，原有逻辑无条件回退到rem_chunk_tokens，绕过SWA预算，在持续压力下同样可能触发OOM。修复后，请求将在资源不足时延迟到下一轮处理。

实现拆解

1. 核心逻辑修复：修改python/sglang/srt/managers/schedule_policy.py中的add_chunked_req方法。在混合SWA模式下，计算_rem_tokens时从int(self.rem_swa_tokens)中减去self.page_size，为alloc_extend预留页面空间。当_rem_tokens <= 0时，如果是混合SWA模式，则直接返回请求对象以延迟处理，而不是回退到rem_chunk_tokens。
2. 测试配套：在test/registered/unit/managers/test_prefill_adder.py中新增两个测试函数。test_add_chunked_req_hybrid_swa_reserves_page_for_alloc_extend验证预留页面逻辑，确保extend_input_len不超过rem_swa_tokens - page_size。test_add_chunked_req_hybrid_swa_defers_when_swa_below_page验证当SWA预算低于页面大小时，请求被延迟而非回退。同时添加辅助函数_build_hybrid_swa_chunked_req以简化测试构建。

关键文件：

• python/sglang/srt/managers/schedule_policy.py（模块 调度策略；类别 source；类型 core-logic；符号 add_chunked_req）: 修复了混合SWA分块预填充的内存溢出问题，是核心逻辑变更所在。
• test/registered/unit/managers/test_prefill_adder.py（模块 预填充测试；类别 test；类型 test-coverage；符号 _build_hybrid_swa_chunked_req, test_add_chunked_req_hybrid_swa_reserves_page_for_alloc_extend, test_add_chunked_req_hybrid_swa_defers_when_swa_below_page）: 新增测试用例验证修复逻辑，确保预留页面和延迟处理行为正确。

关键符号：add_chunked_req

关键源码片段

python/sglang/srt/managers/schedule_policy.py

修复了混合SWA分块预填充的内存溢出问题，是核心逻辑变更所在。

def add_chunked_req(self, req: Req):
    if self.dllm_config is not None:
        _rem_tokens = self._get_dllm_remain_tokens()
    else:
        _rem_tokens = min(self.rem_chunk_tokens, int(self.rem_total_tokens))
        if self.is_hybrid_swa:
            # alloc_extend needs extend_num_tokens + page_size per request,
            # so reserve one page here to avoid OOM
            _rem_tokens = min(
                _rem_tokens, int(self.rem_swa_tokens) - self.page_size
            )
        # The chunked_req must be added to the list; otherwise, it will cause a memory leak.
        # Therefore, in certain cases where _rem_tokens <= 0, it should be replaced with rem_chunk_tokens.
        if _rem_tokens <= 0:
            if self.is_hybrid_swa:
                return req # 混合 SWA 模式下资源不足，延迟请求而非回退
            _rem_tokens = self.rem_chunk_tokens # 非混合 SWA 模式回退到分块令牌预算
​
    truncated = req.extend_input_len > _rem_tokens
    req.set_extend_input_len(min(req.extend_input_len, _rem_tokens))
    req.fill_ids = req.fill_ids[: len(req.prefix_indices) + req.extend_input_len]
    self.can_run_list.append(req)
    self._update_prefill_budget(
        0,
        req.extend_input_len,
        (
            min(req.sampling_params.max_new_tokens, CLIP_MAX_NEW_TOKENS)
            if not truncated
            else 0
        ),
    )
​
    # Return if chunked prefill not finished
    return req if truncated else None

test/registered/unit/managers/test_prefill_adder.py

新增测试用例验证修复逻辑，确保预留页面和延迟处理行为正确。

def test_add_chunked_req_hybrid_swa_reserves_page_for_alloc_extend(self):
    # alloc_extend needs extend_num_tokens + page_size per request. If the
    # scheduler hands out all of rem_swa_tokens, alloc_extend cannot get its
    # extra page and OOMs. With the fix, extend_input_len must cap at
    # rem_swa_tokens - page_size so the page is reserved.
    PAGE_SIZE = 64
    REM_SWA = 100
    adder, req = self._build_hybrid_swa_chunked_req(
        page_size=PAGE_SIZE, rem_swa=REM_SWA
    )
​
    result = adder.add_chunked_req(req)
​
    self.assertIs(result, req) # truncated → chunked prefill continues
    req.set_extend_input_len.assert_called_once()
    new_len = req.set_extend_input_len.call_args.args[0]
    self.assertLessEqual(new_len + PAGE_SIZE, REM_SWA) # 验证预留页面空间
    self.assertEqual(new_len, REM_SWA - PAGE_SIZE) # 验证计算正确
​
def test_add_chunked_req_hybrid_swa_defers_when_swa_below_page(self):
    # When rem_swa_tokens <= page_size there is no room to serve even the
    # reservation, so the chunked req must be deferred (returned unchanged)
    # instead of falling back to rem_chunk_tokens and bypassing SWA budget.
    PAGE_SIZE = 64
    adder, req = self._build_hybrid_swa_chunked_req(
        page_size=PAGE_SIZE, rem_swa=PAGE_SIZE
    )
    original_len = req.extend_input_len
​
    result = adder.add_chunked_req(req)
​
    self.assertIs(result, req) # 请求被延迟，未修改
    req.set_extend_input_len.assert_not_called() # 验证未调用设置长度
    self.assertEqual(req.extend_input_len, original_len) # 长度保持不变
    self.assertEqual(len(adder.can_run_list), 0) # 未加入运行列表

评论区精华

review中仅有gemini-code-assist[bot]的评论，总结了PR的修改内容和目的，指出修改了add_chunked_req方法以防止OOM，并添加了单元测试验证预留逻辑和延迟行为。评论者表示没有反馈，说明变更已被接受。

• 修复混合SWA分块预填充OOM的代码审查 (correctness): 变更被接受，没有进一步反馈。

风险与影响

• 风险：技术风险较低，但需注意：
  • 回归风险：修改了add_chunked_req的核心逻辑，可能影响其他非混合SWA路径或DLLM配置下的行为。但变更范围小，且新增测试覆盖了混合SWA场景，降低了风险。
  • 性能影响：预留页面空间可能略微减少单次可处理的令牌数，但避免了OOM导致的崩溃，整体更稳定。
  • 兼容性：仅影响混合SWA模式，对其他模式无影响。
  • 影响：对系统的影响是修复了内存溢出问题，提升了混合SWA模式下分块预填充的稳定性。用户将不再遇到因OOM导致的服务器崩溃，特别是在高负载或资源紧张时。对团队而言，新增的测试用例增强了相关逻辑的验证，有助于防止未来回归。影响范围限于调度器的预填充处理模块，不涉及其他子系统。
  • 风险标记：核心路径变更, 资源分配逻辑调整

关联脉络

• PR #23252 [Fix] Solve the error lead by _commit_transfer_to_req() when using IntraNode NVLink in PD disaggregation: 同样涉及调度和内存管理相关的bugfix，关注资源分配和错误处理。
• PR #23138 fix: reset empty prefill batch fullness: 涉及调度器中的预填充处理，修复了批次满载标志重置问题，与本PR同属调度模块的bugfix。
• PR #22894 fix(hicache): emit KV events for L2 host cache insertions: 涉及缓存和内存管理，与本PR的KV缓存和资源分配主题相关。