# PR #31538 完整报告

- 仓库：`sgl-project/sglang`
- 标题：[diffusion] support resident layers for DiT
- 合并时间：2026-07-29 16:52
- 原文链接：http://prhub.com.cn/sgl-project/sglang/pull/31538

---

# 执行摘要

- 一句话：支持 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 常驻。新增常驻层参数、激活控制、释放保护等。

```python
# 在 __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