# PR #21499 完整报告

- 仓库：`sgl-project/sglang`
- 标题：Add SWA support for runtime busy memory check
- 合并时间：2026-04-12 15:39
- 原文链接：http://prhub.com.cn/sgl-project/sglang/pull/21499

---

# 执行摘要

本 PR 为混合 SWA（Sliding Window Attention）模型添加了运行时繁忙内存检查支持，修复了此前在 `SWARadixCache` 上调用 `tree_cache.protected_size()` 会抛出 `NotImplementedError` 的问题。通过重构内存检查逻辑，新增双池未缓存令牌计数和特定检查方法，使 SWA 架构能够启用严格内存检查，有助于在高压场景下及早发现内存问题。该变更属于中等重要度的功能扩展，对内存管理和调度器模块有直接价值。

# 功能与动机

**为什么做**：根据 PR body，此前 `self_check_during_busy` 在混合 SWA 模型（具有分离的完整 KV 缓存池和 SWA KV 缓存池）上会失败，因为 `tree_cache.protected_size()` 在 `SWARadixCache` 上未实现。这阻碍了 SWA 模型启用严格内存检查，可能掩盖内存泄漏或溢出风险。

**要解决的问题**：
1. 为 SWA 模型实现可用的繁忙内存检查。
2. 处理流式会话尚未支持检查的情况（跳过）。
3. 保持与非 SWA 模型的兼容性。

# 实现拆解

主要改动集中在 `python/sglang/srt/managers/scheduler_runtime_checker_mixin.py`：

| 模块 | 关键改动 | 说明 |
|------|----------|------|
| 未缓存令牌计算 | 新增 `_get_total_uncached_sizes()` | 统一计算完整池和 SWA 池的未缓存令牌总数，替代原有的两个辅助函数。对于 SWA 池，未缓存令牌 = 分配长度 - max( 缓存保护长度 , swa_evicted_seqlen)。 |
| SWA 特定检查 | 新增 `_self_check_during_busy_swa()` | 为 SWA 模型实现双池守恒不变性检查，分别验证完整池和 SWA 池。 |
| 默认检查提取 | 提取 `_self_check_during_busy_default()` | 从现有逻辑中分离出非 SWA 模型的默认检查，保持代码清晰。 |
| 主检查入口 | 修改 `self_check_during_busy()` | 添加对 `SessionAwareCache` 的早期返回（跳过检查），并根据模型类型分派到 SWA 或默认检查。 |

此外，在两个测试文件中启用了 `SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY` 环境变量，以验证 SWA 场景下的检查功能。

# 评论区精华

由于 review 评论为空，没有公开的技术讨论记录。但从提交历史可以看出实现过程中的设计决策：

1. **初始实现**：`add SWA busy check using flat checker structure` – 基于扁平检查器结构添加 SWA 繁忙检查。
2. **简化重构**：`simplify: merge uncached size helpers into single _get_total_uncached_sizes` – 将多个辅助函数合并为单一方法，减少重复逻辑。
3. **进一步简化**：`simplify batch identity check in _get_total_uncached_sizes` – 优化批次身份检查逻辑，提升代码可读性。

这些提交反映了作者对代码简洁性和可维护性的关注。

# 风险与影响

**技术风险**：
- **回归风险**：修改了核心内存检查路径，可能意外影响非 SWA 模型的检查正确性。但通过保留提取的默认检查方法，并运行现有测试（如 `test_radix_attention.py`、`test_retract_decode.py`）来缓解。
- **性能开销**：新增的双池计算在每次繁忙检查时执行，可能增加少量 CPU 开销，但仅在检查启用时触发，影响有限。
- **覆盖不足**：虽然启用了两个 SWA 测试的环境变量，但缺乏针对新检查逻辑的专项测试，难以确保边界情况（如极端内存压力）下的行为正确。
- **兼容性限制**：当流式会话启用时跳过检查，这是一个已知妥协，可能掩盖会话场景下的内存问题。

**影响范围**：
- **用户**：SWA 模型用户现在可以启用严格内存检查，有助于在并发高压场景下及早发现内存异常，提升系统稳定性。
- **系统**：扩展了内存安全 invariant 对新兴 SWA 架构的支持，保持了内存管理子系统的完整性。
- **团队**：代码结构更清晰，SWA 特定逻辑被分离，便于后续维护和扩展其他缓存类型。

# 关联脉络

与近期历史 PR 的关联揭示了内存检查机制的演进方向：

1. **PR #22562**（Flatten memory checkers into composable per-pool invariant checks）同样修改了 `scheduler_runtime_checker_mixin.py`，将内存检查器重构为可组合的每池检查。本 PR 的 SWA 检查很可能基于该重构后的扁平结构实现，体现了架构一致性。
2. **PR #22577**（Add hisparse staging + decode offload guards to is_fully_idle()）也涉及调度器状态检查逻辑的修复，关注内存和缓存相关的不变量，显示团队对运行时正确性检查的持续投入。

从更广的视角看，本 PR 是 SGLang 对 SWA 等先进注意力机制支持的一部分，与之前添加 SWA 缓存、优化相关内核的 PRs 共同构成了对该架构的完整生态支持。