host_indices 在 HiRadixCache 和 HiCacheController 中的生命周期灵活，容易产生 double-free 问题。该 PR 添加检测机制，使 double-free 时触发断言失败。
该 PR 是典型的防御性编程实践，值得相关模块开发者精读。建议关注：如何使用布尔张量实现分配状态追踪、如何通过测试模拟各种异常路径、以及如何用 benchmark 验证性能无退化。对于正在维护 KV 缓存层或内存池的工程师尤其有参考价值。
Reviewer hzh0425 在审批时建议补充 benchmark 对比数据和长期稳定性测试，以验证改动不会引入异常崩溃。作者在 PR body 中提供了 benchmark 结果（吞吐量无影响），审查者认可后合并。
clear() 中初始化 slot_used 追踪数组python/sglang/srt/mem_cache/pool_host/base.pyclear() 末尾添加 self.slot_used = torch.zeros(self.size, dtype=torch.bool)，用于记录每个 slot 的分配状态。alloc() 中添加 double-alloc 检测select_index 之前，检查 self.slot_used[select_index] 是否有任何 True，若有则触发 AssertionError 并列出重复分配的 slot。通过后将对应 slot 标记为已使用。free() 中添加 double-free 检测AssertionError 并列出未分配的 slot。断言通过后清除 slot_used 标记，再将 slots 追加回 free_slots。test/registered/unit/mem_cache/test_mem_pool_host.pytest_double_alloc（模拟 bookkeeping 损坏）、test_double_free（二次释放）、test_free_unallocated（释放未分配 slot）、test_free_after_clear（clear 后释放旧索引），均验证触发正确的断言消息。register_cpu_ci 注册到 CI 套件 base-a-test-cpu，确保在 CPU CI 中自动运行。| 文件 | 模块 | 状态 | 重要度 |
|---|---|---|---|
python/sglang/srt/mem_cache/pool_host/base.py |
内存池 | modified | 6.08 |
test/registered/unit/mem_cache/test_mem_pool_host.py |
内存池测试 | added | 7.04 |
python/sglang/srt/mem_cache/pool_host/base.py
core-logic
核心逻辑变更，添加了 slot_used 布尔张量并在 alloc/free 中添加断言检测
# 省略导包和类定义，仅展示修改后的核心方法
@synchronized
def clear(self):
# 初始化内存状态和追踪结构
self.mem_state = torch.zeros(
(self.size,), dtype=torch.uint8, device=self.device
)
self.free_slots = torch.arange(self.size, dtype=torch.int64)
# 新增：per-slot 分配标志，True 表示已分配
self.slot_used = torch.zeros(self.size, dtype=torch.bool)
@synchronized
def alloc(self, need_size: int) -> Optional[torch.Tensor]:
assert need_size % self.page_size == 0, "请求大小应为页大小的倍数"
if need_size > self.available_size():
return None
select_index = self.free_slots[:need_size]
self.free_slots = self.free_slots[need_size:]
# 断言：选中的 slot 必须全部未使用
assert not self.slot_used[select_index].any(), (
f"Double-alloc detected: slots already allocated: "
f"{select_index[self.slot_used[select_index]].tolist()}."
)
self.slot_used[select_index] = True
return select_index
@synchronized
def free(self, indices: torch.Tensor) -> int:
indices_cpu = indices.cpu()
# 断言：待释放的 slot 必须全部处于已分配状态
assert self.slot_used[indices_cpu].all(), (
f"Double-free detected: slots not currently allocated: "
f"{indices_cpu[~self.slot_used[indices_cpu]].tolist()}."
)
self.slot_used[indices_cpu] = False
self.free_slots = torch.cat([self.free_slots, indices_cpu])
return len(indices)
test/registered/unit/mem_cache/test_mem_pool_host.py
test-coverage
新增的单元测试文件，覆盖 double-alloc、double-free 等异常路径
import unittest
import torch
from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool
from sglang.srt.mem_cache.memory_pool_host import MHATokenToKVPoolHost
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=2, suite="base-a-test-cpu")
class TestHostKVCache(CustomTestCase):
def setUp(self):
self.page_size = 2
# 构造一个小的 device pool 用于初始化 host pool
self.device_pool = MHATokenToKVPool(
size=self.page_size * 2,
page_size=self.page_size,
dtype=torch.float16,
head_num=2,
head_dim=4,
layer_num=2,
device="cpu",
enable_memory_saver=False,
)
self.host_pool = MHATokenToKVPoolHost(
device_pool=self.device_pool,
host_to_device_ratio=2.0,
host_size=0,
page_size=self.page_size,
layout="layer_first",
pin_memory=False,
device="cpu",
allocator_type="default",
)
def test_double_alloc(self):
indices = self.host_pool.alloc(4)
# 模拟 bookkeeping 损坏：将已分配的 slot 插回 free_slots 头部
leak = torch.tensor([int(indices[0])])
self.host_pool.free_slots = torch.cat([leak, self.host_pool.free_slots])
with self.assertRaises(AssertionError) as ctx:
self.host_pool.alloc(4)
msg = str(ctx.exception)
self.assertIn("Double-alloc", msg)
self.assertIn(f"[{int(leak[0])}]", msg)
def test_double_free(self):
indices = self.host_pool.alloc(4)
self.host_pool.free(indices[:2])
# indices[1] 被二次释放
with self.assertRaises(AssertionError) as ctx:
self.host_pool.free(indices[1:])
msg = str(ctx.exception)
self.assertIn("Double-free", msg)
self.assertIn(f"[{int(indices[1])}]", msg)
def test_free_unallocated(self):
indices = torch.tensor([1])
with self.assertRaises(AssertionError) as ctx:
self.host_pool.free(indices)
msg = str(ctx.exception)
self.assertIn("Double-free", msg)
self.assertIn(f"[{int(indices[0])}]", msg)
def test_free_after_clear(self):
indices = self.host_pool.alloc(4)
self.host_pool.clear()
with self.assertRaises(AssertionError) as ctx:
self.host_pool.free(indices)
msg = str(ctx.exception)
self.assertIn("Double-free", msg)
self.assertIn(str(indices.tolist()), msg)
if __name__ == "__main__":
unittest.main()
Reviewer hzh0425 在审批时建议添加 benchmark 对比数据和长期稳定性测试，以确保改动不会引入异常崩溃。
结论：作者在 PR body 中提供了 benchmark 截图（吞吐量无影响），审查者认可后合并。 · 已解决
该 PR 仅增加断言和布尔标记，风险较低。潜在风险：
1) 新增的 slot_used 布尔张量占用少量额外内存（与 pool size 成正比），但 HostKVCache 规模通常有限，可忽略；
2) 断言检查增加了 alloc/free 路径的计算开销，但 benchmark 显示无显著影响；
3) 如果断言逻辑存在 bug（例如在并发下 @synchronized 是否能保证正确性），但该装饰器已确保序列化访问，正确性较高。
对用户：在调试阶段或生产环境中，若出现 double-free 会立即崩溃并给出明确错误信息，避免静默数据损坏，提升了系统可靠性和可调试性。对系统：增加可忽略的内存和计算开销。对团队：易于定位 HostKVCache 相关的 bug，降低排查成本。
当前没有检测到明确关联的 Issue 链接，后续同步到相关引用后会出现在这里。
参与讨论