执行摘要

• 一句话：支持DiT层常驻GPU加速去噪推理
• 推荐动作：建议diffusion推理用户根据显存余量尝试此功能，特别是需要降低延迟的场景。该PR的设计（延迟激活、force释放、跨模块保护）值得有类似需求的项目参考。如遇到OOM，应减少resident_layers或保持为0。

功能与动机

根据PR描述，现有的--dit-layerwise-offload在显存有限的GPU上有用，但每个去噪步骤都要将所有层从CPU流式传输到GPU。增加--dit-offload-prefetch-size并不能随着显存增加而加速。因此引入--dit-layerwise-resident-layers，让前导层常驻GPU，减少重复传输。

实现拆解

实现分为4个步骤：

1. 参数定义与校验（server_args.py）：在ServerArgs数据类中添加dit_layerwise_resident_layers字段，默认0.0；在CLI解析器中添加--dit-layerwise-resident-layers参数，支持比例或绝对数值；在_validate_offload方法中校验数值非负、截断非整数、以及与offload启用状态的关联警告。

2. Offload管理器扩展（layerwise_offload.py）：在LayerwiseOffloadManager.__init__中添加resident_layers参数并限制范围；新增只读属性holds_residents和_retained_layers，新增方法_activate_residency（由首次前向钩子调用）以延迟激活常驻集；修改prepare_for_next_req使其预取层数为max(prefetch_size, _retained_layers)；修改release_layer增加force参数，默认跳过常驻层（除非force=True）。

3. 组件管理器适配（component_manager.py）：导入is_resident_layerwise_module；在_prefetch_use中跳过持有常驻集的模块的提前预取，避免其他组件阶段占用显存；在finish_request中确保常驻层不跨请求保留（释放）。

4. 测试覆盖（test_layerwise_offload.py）：新增辅助类_MultiBlockModel、_ResidentComponent以及工具函数_patch_fake_device、_resident_manager、_arm_residency；添加两个测试：test_resident_layers_stay_pinned_until_stage_teardown验证常驻层在force释放前保持GPU，test_resident_layers_off_by_default_streams_everything验证默认行为与全流式一致。

关键文件：

• python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py（模块 卸载管理器；类别 source；类型 core-logic；符号 holds_residents, _retained_layers, _activate_residency, release_layer）: 核心实现文件，管理DiT层的CPU offload与GPU常驻。新增常驻层参数、激活控制、释放保护等。
• python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py（模块 层卸载测试；类别 test；类型 test-coverage；符号 _MultiBlockModel, init, _ResidentComponent, _patch_fake_device）: 新增两个单元测试验证常驻层行为，确保功能正确。
• python/sglang/multimodal_gen/runtime/server_args/server_args.py（模块 服务参数；类别 source；类型 core-logic；符号 dit_layerwise_resident_layers, add_cli_args, _validate_offload）: 添加新参数定义、CLI参数和校验逻辑。
• python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py（模块 组件管理器；类别 source；类型 core-logic；符号 _prefetch_use, _finish_use, _should_keep_after_use）: 集成常驻层保护逻辑，避免跨阶段预取和跨请求保留导致OOM。

关键符号：holds_residents, _retained_layers, _activate_residency, release_layer, is_resident_layerwise_module, _prefetch_use, _finish_use, _should_keep_after_use, add_cli_args, _validate_offload

关键源码片段

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

核心实现文件，管理DiT层的CPU offload与GPU常驻。新增常驻层参数、激活控制、释放保护等。

# 在 __init__ 中新增 resident_layers 参数和延迟激活标志
def __init__(
    self,
    model: torch.nn.Module,
    *,
    layers_attr_str: str,
    num_layers: int,
    enabled: bool,
    pin_cpu_memory: bool = True,
    prefetch_size: int = 1,
    resident_layers: int = 0, # 新增：常驻前 N 层
) -> None:
    # ... 原有初始化 ...
    self.resident_layers = min(max(0, int(resident_layers)), self.num_layers)
    self._residency_active = False # 延迟激活标志
​
@property
def holds_residents(self) -> bool:
    """是否持有常驻层（启用且 resident_layers > 0）"""
    return self.enabled and self.resident_layers > 0
​
@property
def _retained_layers(self) -> int:
    """当前跨步骤常驻层数；未激活时返回 0"""
    return self.resident_layers if self._residency_active else 0
​
@torch.compiler.disable
def _activate_residency(self) -> None:
    """在第一个去噪前向时激活常驻集"""
    self._residency_active = True
​
def prepare_for_next_req(self, non_blocking=True):
    """准备下一轮去噪：确保常驻层和预取窗都被预取"""
    num_prefetch_layers = max(self.prefetch_size, self._retained_layers)
    for i in range(num_prefetch_layers):
        self.prefetch_layer(i, non_blocking=non_blocking)
    # 同步流 ...
​
@torch.compiler.disable
def release_layer(self, layer_idx: int, force: bool = False) -> None:
    """释放层。如果常驻层且非强制释放，则跳过"""
    if not self.enabled:
        return
    if not force and layer_idx < self._retained_layers:
        return
    # 原有释放逻辑（替换为占位符）
    if layer_idx in self._gpu_layers:
        self._replace_with_placeholder(layer_idx)
        self._gpu_layers.discard(layer_idx)
​
def is_resident_layerwise_module(module: torch.nn.Module) -> bool:
    """模块是否持有常驻层"""
    if not is_layerwise_offloaded_module(module):
        return False
    return any(
        manager.holds_residents
        for manager in module.layerwise_offload_managers
        if manager is not None
    )

评论区精华

摘要如下：

• mickqian在__init__中询问resident_layers是否应保持之前的值，作者回应这是新参数，默认0无影响。
• mickqian建议合并简化初始化日志，后续提交已处理。

• mickqian提议对is_resident_layerwise_module使用@lru_cache，作者以检查轻量且不在热路径为由拒绝。

• resident_layers默认值 (design): 作者澄清默认值0保持向后兼容，问题解决。

• 初始化日志简化 (style): 作者在后续提交中合并了日志行（日志信息整合为一行含resident_layers）。
• is_resident_layerwise_module使用缓存 (performance): 作者解释后，mickqian未再提异议。

风险与影响

• 风险：技术风险包括：
  • 参数校验中浮点转整数可能因浮点精度导致层数偏差，但代码使用math.floor处理非整数。
  • 延迟激活依赖于前向钩子，若钩子未触发，常驻层不会激活，功能退化为流式（安全降级）。
  • 组件管理器中跳过预取可能导致DiT阶段首次预取延迟，但prepare_for_next_req会在该阶段开始时执行拉取。
  • 与无offload配置一同使用时参数被忽略但打印警告，用户可能误解。
• 影响：影响评估：
  • 用户：使用--dit-layerwise-offload的用户可通过新参数自由调节延迟/显存权衡，默认行为不变。
  • 系统：改动集中在diffusion路径的offload管理层和组件管理器，不影响LLM或其他模块。
  • 团队：设计模式（延迟激活、force释放）可复用于未来类似场景。
  • 风险标记：新增配置参数校验, 常驻层延迟激活依赖第一次前向, 跨模块OOM保护可能引入性能回归, 兼容性警告可能被忽略

关联脉络

• 暂无明显关联 PR