执行摘要
- 一句话：为 Host KV Cache 添加共享内存分配器支持
- 推荐动作：该 PR 值得仔细阅读，尤其是
get_allocator_type 的设计模式（将分配器选择与存储后端配置解耦）以及 ShmHostTensorAllocator 的 fd 追踪方法。讨论中关于 memfd_create 与文件回退、madvise 对比 MAP_POPULATE 的技术选型也值得参考。对于需要理解 HiCache 存储层扩展机制的工程师是必读材料。
功能与动机
在分层缓存 (HiCache) 架构中，为了支持零拷贝数据卸载到动态存储后端（如基于 Unix Domain Socket 的卸载进程），需要在共享内存中分配 host 侧的 KV cache 池。通过共享文件描述符，可以避免数据在 SGLang 进程和存储后端之间的 CPU 内存复制，显著降低二级 KV cache 命中时的检索延迟（测试显示减少约 17.7%）。
实现拆解
- 新增共享内存分配函数：在
mmap_allocator.py 中新增 alloc_shm 函数，通过 memfd_create（优先）或回退到 /dev/shm 文件创建共享内存映射，返回张量及其文件描述符。同时改进 alloc_mmap 使用 madvise(MADV_POPULATE_WRITE) 增强预取。
- 引入共享内存分配器：在
pool_host/common.py 中新增 ShmHostTensorAllocator 类，继承自 HostTensorAllocator，重写 allocate 方法以调用 alloc_shm，并维护文件描述符列表以便在析构时自动关闭。同时新增 get_allocator_type 函数，从 server_args 中解析出分配器类型（支持 shm 后端或通过 dynamic 配置的 allocator 字段）。
- 注册 dummy 存储后端：创建
storage/shm/hicache_shm.py 模块，实现 HiCacheShm 类作为 dummy 后端（因为 shm 是本地分配，不需要实际存储传输），并在 backend_factory.py 中注册 shm 选项。
- 集成到分层缓存构建流程：在
hybrid_pool_assembler.py、hiradix_cache.py 和 decode_kvcache_offload_manager.py 中使用 get_allocator_type 替代直接引用 server_args.hicache_storage_backend，使各缓存池构建时能根据后端动态选择分配器。
- 补充测试：新增
test_mmap_allocator.py 和扩展 test_mem_pool_host.py，覆盖共享内存分配、unlink 行为、hugepage fallback、设备断言以及 ShmHostTensorAllocator 的 fd 管理。
关键文件：
python/sglang/srt/mem_cache/pool_host/common.py（模块 Host 内存池；类别 source；类型 dependency-wiring；符号 ShmHostTensorAllocator, init, fd, mm）: 核心文件：新增 ShmHostTensorAllocator 和 get_allocator_type 函数，是分配器选择逻辑的心脏。python/sglang/srt/mem_cache/storage/mmap/mmap_allocator.py（模块 内存分配器；类别 source；类型 rename-or-move；符号 alloc_shm）: 底层分配函数：新增 alloc_shm 函数并优化 alloc_mmap 的预取策略。test/registered/unit/mem_cache/test_mmap_allocator.py（模块 测试；类别 test；类型 test-coverage；符号 TestMmapAllocator, test_alloc_mmap, test_alloc_shm, test_shm_host_tensor_allocator）: 新增共享内存分配器单元测试，覆盖基本功能、unlink、hugepage fallback 和设备断言。python/sglang/srt/mem_cache/storage/shm/hicache_shm.py（模块 HiCache 存储；类别 source；类型 core-logic；符号 HiCacheShm, init, get, batch_get）: 新增 dummy 存储后端 HiCacheShm，用于注册 shm 分配器类型以兼容现有后端工厂。python/sglang/srt/mem_cache/storage/mmap/__init__.py（模块 内存分配器；类别 source；类型 dependency-wiring）: 新模块初始化文件，导出 alloc_mmap 和 alloc_shm 供外部调用。python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py（模块 缓存池构建；类别 source；类型 core-logic；符号 _get_allocator_type）: 使用 get_allocator_type 重构分配器选择，使各缓存池能根据后端动态切换。
关键符号：alloc_shm, ShmHostTensorAllocator.allocate, get_allocator_type, HiCacheShm.init
关键源码片段
python/sglang/srt/mem_cache/pool_host/common.py
核心文件：新增 ShmHostTensorAllocator 和 get_allocator_type 函数，是分配器选择逻辑的心脏。
class ShmHostTensorAllocator(HostTensorAllocator):
def __init__(self):
super().__init__()
self.fds = [] # 记录所有分配的文件描述符，以便析构时关闭
self.mms = []
@property
def fd(self):
# 返回第一个分配的文件描述符
return self.fds[0] if self.fds else None
@property
def mm(self):
return self.mms[0] if self.mms else None
def allocate(self, dims: tuple, dtype: torch.dtype, device: str) -> torch.Tensor:
# 确保设备为 CPU，共享内存只支持 CPU 分配
assert device == "cpu", f"ShmHostTensorAllocator only supports CPU allocations; got device={device!r}"
self.dtype = dtype
self.dims = dims
from sglang.srt.mem_cache.storage.mmap import alloc_shm
tensor, fd, mm = alloc_shm(dims, dtype) # 分配共享内存并返回 fd 和 mmap 对象
self.fds.append(fd)
self.mms.append(mm)
return tensor
def __del__(self):
# 析构时关闭所有记录的文件描述符，避免泄漏
for fd in getattr(self, "fds", []):
if fd is not None:
try:
os.close(fd)
except OSError:
pass
self.fds = []
def get_allocator_type(server_args) -> str:
# 从 server_args 中推导出分配器类型字符串
backend = getattr(server_args, "hicache_storage_backend", None)
if backend == "shm":
return "shm"
if backend == "dynamic":
extra_config_str = getattr(server_args, "hicache_storage_backend_extra_config", None)
if extra_config_str:
try:
config = json.loads(extra_config_str)
if config.get("allocator") == "shm":
return "shm"
except Exception:
pass
return backend or "default"
python/sglang/srt/mem_cache/storage/mmap/mmap_allocator.py
底层分配函数：新增 alloc_shm 函数并优化 alloc_mmap 的预取策略。
def alloc_shm(dims: tuple, dtype: torch.dtype) -> tuple[torch.Tensor, int, mmap.mmap]:
# 通过共享内存分配张量，返回 (tensor, fd, mm) 元组
n_bytes = math.prod(dims) * torch.empty([], dtype=dtype).element_size()
hugepage_size = (envs.SGLANG_HUGEPAGE_SIZE.get() or "").strip().upper()
if hugepage_size != "":
logger.warning("Hugepages are not supported with SHM allocator. Falling back to plain page-size mmap.")
page_size = mmap.PAGESIZE
alloc_bytes = math.ceil(n_bytes / page_size) * page_size
# 优先使用 memfd_create（内核 3.17+）
fd = None
try:
fd = os.memfd_create(
f"sglang_host_pool_{uuid.uuid4().hex}",
flags=getattr(os, "MFD_CLOEXEC", 1),
)
except (AttributeError, OSError):
# 回退到 /dev/shm 中的临时文件
shm_path = f"/dev/shm/sglang_host_pool_{uuid.uuid4().hex}.mmap"
try:
fd = os.open(shm_path, os.O_CREAT | os.O_RDWR | os.O_TRUNC, 0o600)
os.unlink(shm_path) # 立即 unlink，OS 回收内存
except Exception as e:
raise OSError(f"Failed to create shm file: {e}")
os.ftruncate(fd, alloc_bytes)
mm = mmap.mmap(fd, alloc_bytes, flags=mmap.MAP_SHARED | _MAP_POPULATE, prot=mmap.PROT_READ | mmap.PROT_WRITE)
# MADV_POPULATE_WRITE 改进了预取行为（内核 5.14+）
try:
mm.madvise(_MADV_POPULATE_WRITE)
except OSError:
# 旧内核回退到 MAP_POPULATE
pass
tensor = torch.frombuffer(mm, dtype=dtype, count=math.prod(dims)).reshape(dims)
return tensor, fd, mm
评论区精华
风险与影响
- 风险：
- 共享内存泄漏：/dev/shm 中的文件在创建后立即 unlink，由 OS 自动回收；但若 fd 在 ShmHostTensorAllocator 析构前丢失，可能泄漏 fd（已通过 del 清理）。低风险。
- Hugepage 回退静默：对 hugepage 请求，alloc_shm 会发出警告并回退到普通页面，可能影响性能（测试已覆盖）。
- 与 Hybrid Pool 集成验证不充分：虽然声称兼容，但多池场景（如 Full + Indexer）下 fd 管理可能有未释放风险，需要更多端到端测试。
- 旧内核兼容性：MADV_POPULATE_WRITE 在 <5.14 内核上回退到 MAP_POPULATE，不影响功能但预取弱化。
- 影响：用户可通过 --hicache-storage-backend shm 启用共享内存分配，但依赖 Linux memfd_create（内核 3.17+）或 /dev/shm 存在。零拷贝 offloading 将基于此特性实现，是该功能链路上的基础模块。代码结构重构（mmap_allocator 搬迁至 storage/mmap/，新增 storage/shm/），团队后续 offloading PR 需基于此接口。
- 风险标记：共享内存泄漏, hugepage 静默回退, 混合池兼容需验证, 旧内核预取降级
关联脉络
参与讨论