执行摘要

• 一句话：按 cgroup 上限规划 pinned 主机内存预算
• 推荐动作：值得精读。该 PR 展示了"从容器真实限制出发做资源预算"的设计思路，HostPinBudget 的热度排序（字节 × 步数）比简单的 DiT 优先更符合传输量实际，且把回退代价量化进 docstring 和日志，是资源治理的好范例。关注 host_memory_budget.py 的 cgroup 解析与 layerwise_offload.py 的排序注入点。

功能与动机

PR body 指出：offloaded weights 放在 pinned host memory 中，内核既不能 swap 也不能回收，而 pin_cpu_memory 只是一个全局布尔值，运行时在提交前从未查看剩余主机内存。psutil.virtual_memory() 读 /proc/meminfo 只反映宿主机而看不到容器限制，实测一台租用 4x RTX 4090 机器 psutil 报告 2015.7 GiB 而 cgroup 上限为 1117.2 GiB，900 GiB 的虚报正好会把容器推向 OOM-kill。因此必须直接读取 cgroup 的 memory cap 来规划 pin 预算。

实现拆解

1. 新增 python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_budget.py：实现 _read_int（安全读取文件并处理 max 关键字）、cgroup_memory_limit_bytes（按 v2 优先、v1 回退读取 (cap, usage)，用 _UNLIMITED_ABOVE 哨兵识别 v1 的近乎 2^63 无限上限）、host_memory_available_bytes（取 psutil.virtual_memory().available 与 max(0, limit - usage) 的较小值）、module_weight_bytes（统计模块参数与 buffer 字节数）与 pin_benefit_bytes（权重字节 × 每请求使用次数）。
2. 核心类 HostPinBudget：构造时计算 available_bytes 与 reserve_bytes = max(5%, 2 GiB)，spendable_bytes 为扣除储备和已承诺后的余量；request(component_name, weight_bytes) 在余量内承诺并返回 True，否则记日志并返回 False。日志与 docstring 明确量化了 pageable 回退的代价（text encoding 阶段 3.04 s -> 6.26 s，约 2x；H2D 带宽 pinned 13.09 GB/s vs pageable 8.84 GB/s，且 copy_(non_blocking=True) 对 pageable 源会静默变同步）。
3. 修改 layerwise_offload.py：configure_layerwise_offload 签名增加 pin_budget 与 component_name 参数，逐 layer 用 pin_budget.request(...) 决定该 manager 是否 pin；新增 _default_num_inference_steps 从 pipeline 配置类读取默认采样步数，新增 _h2d_bytes_a_pin_would_save 计算每个组件的 pin 收益（DiT 按步数多次传输）；configure_layerwise_offload_modules 每个配置轮次构建一个 budget，并按 pin_benefit_bytes 降序 offer 给各组件，拒绝的组件回退为 pageable。
4. 测试配套：新增 test_host_memory_budget.py（15 个单测，覆盖 v2/v1 路径、无 cgroup、两种 unlimited 写法、cap 与 free 取小、储备及下限、热组件优先耗尽、体积估算去重）；test_layerwise_offload.py 为 _TestServerArgs 补充 pipeline_class_name=None 字段，避免 6 个既有测试因新读取逻辑而失败。

关键文件：

• python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_budget.py（模块 内存预算；类别 source；类型 dependency-wiring；符号 _read_int, cgroup_memory_limit_bytes, host_memory_available_bytes, HostPinBudget）: 新增的核心模块：从 cgroup v1/v2 读取容器内存上限，结合 psutil 可用内存计算可提交预算，并提供 HostPinBudget 按请求分配 pin 额度。是本次变更的入口和策略载体。
• python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py（模块 层间卸载；类别 source；类型 core-logic；符号 configure_layerwise_offload, _default_num_inference_steps, _h2d_bytes_a_pin_would_save）: 修改核心调度逻辑：在 configure_layerwise_offload 中按 pin_budget 逐层决策是否 pin，并在 configure_layerwise_offload_modules 中按 pin_benefit_bytes 排序分配预算。
• python/sglang/multimodal_gen/test/unit/test_host_memory_budget.py（模块 单元测试；类别 test；类型 test-coverage；符号 _point_at, TestCgroupLimit, test_v2_cap_is_read, test_v1_cap_is_read_when_v2_is_absent）: 新增 15 个单元测试，覆盖 cgroup v1/v2 解析、无 cgroup、unlimited 哨兵、cap 与 free 取小、储备、预算耗尽与存储去重等关键路径，是验证预算逻辑的主要手段。
• python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py（模块 单元测试；类别 test；类型 test-coverage）: 为 _TestServerArgs 补充 pipeline_class_name 字段，修复因新增默认步数读取而破坏的 6 个既有测试。

关键符号：configure_layerwise_offload, _default_num_inference_steps, _h2d_bytes_a_pin_would_save, cgroup_memory_limit_bytes, host_memory_available_bytes, HostPinBudget.request, HostPinBudget.spendable_bytes, module_weight_bytes, pin_benefit_bytes

关键源码片段

python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_budget.py

新增的核心模块：从 cgroup v1/v2 读取容器内存上限，结合 psutil 可用内存计算可提交预算，并提供 HostPinBudget 按请求分配 pin 额度。是本次变更的入口和策略载体。

# host_memory_budget.py
# 关键点：pin 过的页内核既不能 swap 也不能回收，
# 所以提交前必须按容器 cgroup 上限而非宿主机总内存规划。
​
import psutil
​
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
​
logger = init_logger(__name__)
​
GIB_BYTES = 1024**3
​
# cgroup v2 与 v1 的 limit/usage 文件路径，v2 优先
_CGROUP_V2 = ("/sys/fs/cgroup/memory.max", "/sys/fs/cgroup/memory.current")
_CGROUP_V1 = (
    "/sys/fs/cgroup/memory/memory.limit_in_bytes",
    "/sys/fs/cgroup/memory/memory.usage_in_bytes",
)
​
# v1 的 " 无限 " 上限用接近 2**63 的哨兵值表示，超过该值视为无 cap
_UNLIMITED_ABOVE = 1 << 62
​
# 保留 5%（至少 2 GiB）不参与 pin，留给激活、暂存缓冲区等中间分配
HOST_RESERVE_FRACTION = 0.05
MIN_HOST_RESERVE_BYTES = 2 * GIB_BYTES
​
​
def _read_int(path: str) -> int | None:
    """读取文件中的整数，文件缺失、"max" 或非法内容都返回 None。"""
    try:
        with open(path) as handle:
            text = handle.read().strip()
    except OSError:
        return None
    if text == "max": # cgroup v2 的 "max" 表示无限制
        return None
    try:
        return int(text)
    except ValueError:
        return None
​
​
def cgroup_memory_limit_bytes() -> tuple[int, int] | None:
    """返回 (cap, usage)，无 cgroup 或无限时返回 None。"""
    for limit_path, usage_path in (_CGROUP_V2, _CGROUP_V1):
        limit = _read_int(limit_path)
        # 无限制或哨兵过大时跳过，尝试下一个版本
        if limit is None or limit >= _UNLIMITED_ABOVE:
            continue
        usage = _read_int(usage_path) or 0
        return limit, usage
    return None
​
​
def host_memory_available_bytes() -> int:
    """进程还能安全提交的字节数：宿主机可用与 cgroup 余量的较小值。"""
    available = int(psutil.virtual_memory().available)
    capped = cgroup_memory_limit_bytes()
    if capped is None:
        return available
    limit, usage = capped
    return min(available, max(0, limit - usage))
​
​
class HostPinBudget:
    """按热度分配 pin 额度：热组件先问，冷组件在预算耗尽后回退 pageable。"""
​
    def __init__(self, available_bytes: int | None = None) -> None:
        if available_bytes is None:
            available_bytes = host_memory_available_bytes()
        self.available_bytes = available_bytes
        # 储备取 5% 与 2 GiB 中的较大者，防止运行时后续分配无内存可用
        self.reserve_bytes = max(
            int(available_bytes * HOST_RESERVE_FRACTION), MIN_HOST_RESERVE_BYTES
        )
        self.committed_bytes = 0
​
    @property
    def spendable_bytes(self) -> int:
        return max(0, self.available_bytes - self.reserve_bytes - self.committed_bytes)
​
    def request(self, *, component_name: str, weight_bytes: int) -> bool:
        """若 weight_bytes 可被 pin 则承诺并返回 True，否则记日志返回 False。
​
        对热组件也是硬 cap：超额授予不会减小内存占用，只会把失败推迟到
        pin 分配本身或 swap 阶段。优先级靠"先问"体现，而不是靠超发。
        """
        if weight_bytes <= 0:
            return True
        if weight_bytes <= self.spendable_bytes:
            self.committed_bytes += weight_bytes
            return True
        logger.info(
            "Host pin budget: refusing %s (%d bytes), only %d bytes left",
            component_name,
            weight_bytes,
            self.spendable_bytes,
        )
        return False

python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py

修改核心调度逻辑：在 configure_layerwise_offload 中按 pin_budget 逐层决策是否 pin，并在 configure_layerwise_offload_modules 中按 pin_benefit_bytes 排序分配预算。

# layerwise_offload.py 中与 pin 预算相关的改造片段
​
    def configure_layerwise_offload(
        self,
        server_args: ServerArgs,
        *,
        pin_budget: HostPinBudget | None = None,
        component_name: str | None = None,
    ):
        self.layerwise_offload_managers = []
        named_modules = dict(self.named_modules())
        configured_layer_names = []
        # 既有 DiT 调优旋钮（prefetch/resident）仍只作用于 DiT 组
        dit_tuning_enabled = self.layerwise_offload_dit_group_enabled
        for layer_name in self.layer_names:
            module_list = named_modules.get(layer_name)
            if not isinstance(module_list, (torch.nn.ModuleList, torch.nn.Sequential)):
                continue
            if len(module_list) == 0:
                continue
​
            num_layers = len(module_list)
            # ... 原有 prefetch_size / resident_layers 计算 ...
​
            # pin 这些权重是拷贝流超前于计算的前提，但 pinned 页不可回收，
            # 所以只有当预算允许时才 pin；预算会被 request 逐层扣除。
            pin_cpu_memory = server_args.pin_cpu_memory
            if pin_cpu_memory and pin_budget is not None:
                pin_cpu_memory = pin_budget.request(
                    component_name=f"{component_name or type(self).__name__}.{layer_name}",
                    weight_bytes=module_weight_bytes(module_list),
                )
            manager = LayerwiseOffloadManager(
                model=self,
                layers_attr_str=layer_name,
                num_layers=num_layers,
                enabled=True,
                pin_cpu_memory=pin_cpu_memory,
                prefetch_size=prefetch_size,
                resident_layers=resident_layers,
                initialize=False,
            )
            # ... 后续注册 manager ...
​
    @staticmethod
    def _default_num_inference_steps(server_args) -> int:
        """读取 pipeline 配置类里的默认推理步数，用于估算每请求传输次数。"""
        from sglang.multimodal_gen.registry import get_pipeline_config_classes
​
        pipeline_class_name = server_args.pipeline_class_name
        if not pipeline_class_name:
            return 1
        config_classes = get_pipeline_config_classes(pipeline_class_name)
        if config_classes is None:
            return 1
        return max(1, int(config_classes[1]().num_inference_steps))
​
    def _h2d_bytes_a_pin_would_save(self, name: str) -> int:
        """pin 某组件能省下的传输字节：权重字节 × 每请求使用次数。
​
        用乘积而非"是否为 DiT"排序，是因为 few-step 模型下 1 GB DiT 步进
        4 次移动 4 GB/请求，而 20 GB 文本编码器一次移动 20 GB，
        只看 DiT 布尔值会把预算给到更需要的一方。
        """
        module = self.modules[name]
        if not isinstance(module, LayerwiseOffloadableModuleMixin):
            return 0
        return pin_benefit_bytes(
            weight_bytes=module_weight_bytes(module),
            uses_per_request=(
                default_steps if module.layerwise_offload_dit_group_enabled else 1
            ),
        )

python/sglang/multimodal_gen/test/unit/test_host_memory_budget.py

新增 15 个单元测试，覆盖 cgroup v1/v2 解析、无 cgroup、unlimited 哨兵、cap 与 free 取小、储备、预算耗尽与存储去重等关键路径，是验证预算逻辑的主要手段。

# test_host_memory_budget.py
# 用 monkeypatch 把 cgroup 路径重定向到 tmp_path，避免触碰真实系统文件。
​
def _point_at(monkeypatch, tmp_path, *, v2=None, v1=None):
    """把 _CGROUP_V2 / _CGROUP_V1 指向临时文件，构造 v2/v1 场景。"""
​
    def write(name, value):
        path = tmp_path / name
        path.write_text(str(value))
        return str(path)
​
    missing = str(tmp_path / "absent")
    v2_paths = (
        (write("memory.max", v2[0]), write("memory.current", v2[1]))
        if v2
        else (missing, missing)
    )
    v1_paths = (
        (write("limit_in_bytes", v1[0]), write("usage_in_bytes", v1[1]))
        if v1
        else (missing, missing)
    )
    monkeypatch.setattr(host_memory_budget, "_CGROUP_V2", v2_paths)
    monkeypatch.setattr(host_memory_budget, "_CGROUP_V1", v1_paths)
​
​
class TestCgroupLimit:
    def test_v2_cap_is_read(self, monkeypatch, tmp_path):
        _point_at(monkeypatch, tmp_path, v2=(32 * GIB_BYTES, 4 * GIB_BYTES))
        assert cgroup_memory_limit_bytes() == (32 * GIB_BYTES, 4 * GIB_BYTES)
​
    def test_v1_cap_is_read_when_v2_is_absent(self, monkeypatch, tmp_path):
        _point_at(monkeypatch, tmp_path, v1=(64 * GIB_BYTES, 8 * GIB_BYTES))
        assert cgroup_memory_limit_bytes() == (64 * GIB_BYTES, 8 * GIB_BYTES)
​
    def test_no_cgroup_reports_uncapped(self, monkeypatch, tmp_path):
        _point_at(monkeypatch, tmp_path)
        assert cgroup_memory_limit_bytes() is None
​
    def test_v2_max_keyword_is_uncapped(self, monkeypatch, tmp_path):
        _point_at(monkeypatch, tmp_path, v2=("max", 4 * GIB_BYTES))
        assert cgroup_memory_limit_bytes() is None
​
    def test_v1_sentinel_is_uncapped(self, monkeypatch, tmp_path):
        # 无限制的 v1 cgroup 报告接近 2**63 的数字而非 "max"
        _point_at(monkeypatch, tmp_path, v1=(2**63 - 4096, 8 * GIB_BYTES))
        assert cgroup_memory_limit_bytes() is None
​
    def test_the_cap_wins_over_what_the_kernel_reports_free(
        self, monkeypatch, tmp_path
    ):
        # 租用机上实测场景：psutil 看到整台机器，cgroup 是真实上限
        _point_at(monkeypatch, tmp_path, v2=(32 * GIB_BYTES, 8 * GIB_BYTES))
        monkeypatch.setattr(
            host_memory_budget.psutil,
            "virtual_memory",
            lambda: type("VM", (), {"available": 900 * GIB_BYTES})(),
        )
        assert host_memory_available_bytes() == 24 * GIB_BYTES

评论区精华

该 PR 无 review 评论线程，提交信息体现了两次设计迭代：第一次提交按"是否是 DiT"分配预算，第二次提交改为按"字节 × 步数"（pin_benefit_bytes）排序，理由是 few-step 模型下小 DiT 步进多次移动 4 GB/请求，而 20 GB 文本编码器一次移动 20 GB，布尔判断会反向分配。第三次提交修复测试桩缺失 pipeline_class_name 字段导致 6 个既有测试失败的问题。

• 暂无高价值评论线程

风险与影响

• 风险：主要风险集中在回退路径：当组件因预算不足改为 pageable 后，copy_(non_blocking=True) 会退化为同步拷贝，实测 text encoding 阶段耗时近 2 倍（3.04 s -> 6.26 s），端到端约慢 23%，对延迟敏感的多租户场景可能放大排队效应。cgroup 解析逻辑对 v1/v2 的路径和哨兵处理依赖内核行为，在系统监管缺失（无 cgroup 挂载或权限受限）时回退为宿主机 free 值，无法真正防御容器超卖；HostPinBudget 为每个配置轮次新建，若配置轮次间共享进程状态则可能重复承诺，当前实现按模块内单次构建规避了该问题。
• 影响：影响 layerwise offload 在容器内运行的 diffusion 服务：之前可能因 pinned 内存无界提交触发 OOM-kill，现在预算在 cgroup cap 内硬性约束；热组件优先 pin，冷组件（如大文本编码器）可能被降级为 pageable，带来 stage 级 2x 延迟但避免进程死亡。对非容器宿主机或内存充裕场景行为不变，属于防御性增强。团队后续可基于该机制暴露 pin 预算与回退日志，进一步做容量规划与告警。
• 风险标记：容器内存上限误判风险, pageable 回退导致同步拷贝变慢, cgroup 解析依赖内核行为, 主机内存预算为硬 cap 可能误伤热组件, 测试桩依赖新增字段

关联脉络

• PR #35668 placeholders: 同作者在同一 memory_managers 目录下的后续配套（placeholder，占位说明）
• PR #35618 [diffusion] UX: report where a component's weights are: 同为 layerwise offload 组件驻留的可见性改进，本 PR 的日志报告与预算拒绝日志互补
• PR #35626 [diffusion] fix: keep large vocab tables in host memory under layerwise offload: 处理 layerwise offload 下大表驻留内存的策略，与本 PR 的主机内存规划同属内存驻留治理