AkiYang推理工程手记

小 Tensor 为什么保活大内存SGLang PR #39120 的 View 缓存修复

版本与源码依据
系列阅读 · SGLang PR 精读展开篇目
本页目录

一个逻辑大小只有 0.8 MiB 的 Tensor,可能让约 7.8 MiB 的底层存储一直无法释放。原因是它来自一个更大的 batch Tensor 的切分:小 Tensor 保留了对整块 Storage 的引用。

SGLang PR #39120 修复了多模态 Embedding Cache 中的这一问题。旧缓存按可见元素大小记账,却可能保活更大的资源;补丁在需要时把入库的 View 复制成独立存储。下面用一个最小例子解释引用、内存计数和复制时机。

本文固定到 sgl-project/sglang@24b6c1c7f5ba6c95a96acc9e77b4f5efb23eff97。PyTorch View / Storage 语义参考官方 Tensor Views、Storage 与 Tensor.expand 文档。

一、小 View 怎样保活整个 Batch

先不看 SGLang。假设一次批量多模态编码处理 10 个 item,每个 item 产生 100 个 token,hidden size 为 4096,dtype 为 BF16。

整个 Batch Embedding 是:

[1000, 4096]

1000 × 4096 × 2 bytes
≈ 7.81 MiB

现在把它拆成 10 份:

item_embeddings = torch.split(
    batch_embedding,
    [100, 100, 100, 100, 100, 100, 100, 100, 100, 100],
    dim=0,
)

每个 item 的逻辑 shape 是 [100, 4096],所以:

payload
= 100 × 4096 × 2 bytes
≈ 0.78 MiB

如果只看 tensor.numel() * tensor.element_size(),你自然会认为这个 Entry 大约只占 0.78 MiB。

但 PyTorch Tensor 不只有 shape 和 dtype。更完整地看,一个 strided Tensor 还关联着 Storage、stride、storage offset 等元数据;承载底层字节的是 Storage。PyTorch 官方文档也明确说明:View 与 base Tensor 可以共享同一个底层 Storage,而 split() 属于 View operation。

所以刚才的结构更接近:

                    one Batch Storage
┌──────────────────────────────────────────────────┐
│ img0 │ img1 │ img2 │ img3 │ ... │ img8 │ img9 │
└──────────────────────────────────────────────────┘
   ↑      ↑      ↑      ↑              ↑      ↑
 view0  view1  view2  view3           view8  view9

创建这些 View 本来是好事:没有额外 copy,split() 很便宜。

问题出现在:一个短生命周期计算流里的 View 被放进了长期存在的 Cache。

假设 Batch 计算结束后,只剩 view3 还在 Cache:

batch_embedding 对象不再需要
        │
        ▼
view3 仍然存活
        │
        ▼
它仍引用原来的 Batch Storage
        │
        ▼
整块 Batch Storage 不能释放

于是可能出现:

Cache Entry 逻辑 payload:0.78 MiB
实际因它而继续存活的 Storage:7.81 MiB

10 个 View 若来自同一 Batch,会共享同一块 7.81 MiB Storage。统计底层占用时,这块存储只计一次。

危险的是 Cache 长时间运行之后:

Batch A → 只剩一个小 View 被缓存 → pin 住 Batch A Storage
Batch B → 只剩一个小 View 被缓存 → pin 住 Batch B Storage
Batch C → 只剩一个小 View 被缓存 → pin 住 Batch C Storage
...

这时账面上的多个小 Entry,就可能分别延长多个历史 Batch Storage 的生命周期。

这属于 unintended retention:缓存引用延长了 Storage 的生命周期,使整块 Batch 存储持续驻留。

二、逻辑大小与底层 Storage

缓存记账涉及两个不同的字节数。

第一个:

tensor.numel() * tensor.element_size()

它描述的是 Tensor 逻辑 view 中有多少元素数据。PyTorch 的 Tensor.nbytes 也是按 numel() * element_size() 定义。

第二个:

tensor.untyped_storage().nbytes()

untyped_storage() 返回 Tensor 的底层 UntypedStorage;Storage 是保存实际字节的数据容器。

对一个普通独立 Tensor:payload bytes ≈ storage bytes

但对一个从大 Batch 切出来的 View:

View shape: [100, 4096]
Logical payload: 0.78 MiB

Underlying Storage:
[1000, 4096]
Storage bytes: 7.81 MiB

于是:

Logical Tensor Size
        ≠
Physical Backing Storage Size

这也解释了旧版 Cache 为什么会低估 retained memory。修复前的 helper 是:

def _get_tensor_size(embedding):
    return embedding.element_size() * embedding.numel()

也就是默认:Cache Cost = Tensor Payload

但 Cache 需要约束的是:只要这个 Entry 还活着,哪些底层资源必须继续活着?

reshape() 可能返回 View,也可能创建新 Tensor。资源核算因此检查最终 Storage 身份与字节数,确认结果是否仍与原 Tensor 共享存储。

三、缓存记账怎样遗漏真实占用

回到固定 merge commit 的 SGLang 源码。

mm_schedule.py 在批量多模态编码得到一个整体 Tensor 时,会先:

all_miss_embedding = data_embedding_func(miss_items)

如果返回值不是 per-item list,而是一个整体 Tensor,就会执行:

split_embeddings = torch.split(
    all_miss_embedding,
    token_counts,
    dim=0,
)

固定版本源码甚至直接写了注释:一个 torch.split View 会让整个 concatenated buffer 在任意单个 item 仍被缓存时继续存活。对应源码可见 mm_schedule.py 和 per-item 路径。

数据流是:

flowchart TD
    B["Batched multimodal embedding · large shared Storage"]
    S["torch.split()"]
    V1["small View A · small payload"]
    V2["small View B · small payload"]
    C["Long-lived Multimodal Cache"]
    A["Old accounting · numel × element_size"]
    R["Reality · View may retain whole Batch Storage"]

    B --> S
    S --> V1
    S --> V2
    V1 --> C
    V2 --> C
    C --> A
    V1 -. shares .-> R
    V2 -. shares .-> R

torch.split() 通过 View 复用存储,适合短期计算;进入长期缓存后,引用关系延长,整块 Storage 的驻留成本也随之延长。

Root Cause 是:共享 Storage 的短生命周期 View,被带进了按 Entry 独立记账的长期 Cache,而 accounting model 没有反映这种 ownership / lifetime。

因此这个 PR 的核心 invariant 可以写成:

缓存预算需要覆盖 Entry 实际保活的资源。按 Entry 累加记账时,入库数据的存储占用也应与账面大小对应。

这条预算规则涉及两类关系:

Accounting
Lifetime

一个缓存条目“拥有”的不只是一个 Python Tensor 对象,它还会延长与该 Tensor 关联资源的 lifetime。

四、在入库前分离不必要的共享存储

PR #39120 的核心改动位于 multimodal_cache.py。

第一步,把 _get_tensor_size() 改成:

def _get_tensor_size(embedding):
    return embedding.untyped_storage().nbytes()

这让已经入库的 Entry 按真实 backing storage 字节数记账。

入库时的处理在 set() 中:

tensor = embedding.embedding
storage_size = _get_tensor_size(tensor)
data_size = min(
    storage_size,
    tensor.element_size() * tensor.numel(),
)

如果 Entry 能被 Cache 接纳,并且:storage_size > data_size

才做:

embedding = replace(
    embedding,
    embedding=tensor.clone(),
)

最后重新按 clone 后的实际 Storage 大小计入 current_size。

为什么不直接让每个 View 都按完整 storage_size 记账?

因为多个 Entry 可能共享同一个 Storage。假设:

View A ─┐
View B ─┼──→ same 100 MiB Storage
View C ─┘

如果每个 Entry 都独立记 100 MiB,Cache 账面可能变成 300 MiB,但底层其实只有一块 100 MiB Storage。要彻底做 shared-storage accounting,需要维护 Storage identity、引用关系和 per-storage accounting,复杂度会明显增加。

PR 选择的是另一条更简单的路径:

既然 Cache 按 Entry 独立记账,那就让长期缓存的 oversized View 拥有自己的独立 Storage。

也就是:

Before

Batch Storage
████████████████████████████████████
        ↑
      small View

After clone

Batch Storage
████████████████████████████████████

Cached Entry Storage
████

这是一个很典型的工程取舍:通过改变 ownership,让 accounting model 重新变得简单而正确。

从补丁的位置看,复制成本被放在入库时承担:Entry 被接纳且 storage_size > data_size 时,才创建独立存储。

为什么只 clone 一部分 Tensor?

PR 没有无脑 clone() 所有 Entry。

  • 普通独立 Tensor:storage_size == payload_size,不 clone;
  • duplicate key:在更早的位置直接返回,不 clone;
  • Cache 无论如何都无法接纳的 Entry:在 clone 前返回 False,不做无意义 copy;
  • oversized backing Storage 的被接纳 Entry:clone 成独立 Storage。

短期计算继续共享 View,长期缓存则在需要时取得独立 Storage。

min(storage_size, payload_size) 为什么重要?

因为还有反方向的情况:逻辑 payload 可以比真实 Storage 大。

典型例子是 expand()。PyTorch 官方文档明确说明 expand() 不分配新内存,而是通过 stride=0 创建 View。

因此可能有:payload_size > storage_size

此时并不存在“小 Tensor pin 大 Storage”的问题。用 min(storage_size, payload_size) 可以避免把这种 Tensor 错判成需要 compact 的对象。

整个修复流程可以浓缩成:

flowchart TD
    T["Tensor candidate"]
    S["storage_size"]
    P["payload_size"]
    M["target cost = min(storage, payload)"]
    F{"Cache can admit?"}
    C{"storage > target cost?"}
    CL["clone to independent Storage"]
    KEEP["keep original Tensor"]
    ADD["recompute actual storage · and cache"]

    T --> S
    T --> P
    S --> M
    P --> M
    M --> F
    F -- No --> X["Reject · no copy"]
    F -- Yes --> C
    C -- Yes --> CL --> ADD
    C -- No --> KEEP --> ADD

五、验证修复与兼容行为

PR 报告的 CPU Smoke Checks 覆盖普通切片、detached、strided、inference-mode、expanded Tensor 和 LRU 记账,Base / Extra CI 最终通过。验证重点是数值保留、存储独立性与预算更新;GPU 性能和模型精度未单独测量。

修复后的 regression test 也从:

test_tensor_cache_entries_share_storage()

变成:

test_tensor_cache_entries_own_storage()

并直接断言:

emb.untyped_storage().nbytes()
==
emb.numel() * emb.element_size()

也就是说,测试不再只验证 Tensor 数值,而是把缓存 Entry 的 Storage ownership contract也写进回归测试。对应固定版本测试见 test_mm_chunked_embedding_unit.py。

这个断言把资源生命周期也纳入回归检查:数值正确的 Tensor,仍可能引用过大的底层 Storage。对于本 PR 的缓存切分案例,测试同时检查这两件事。

六、把内存预算落到资源所有权

这次修复把三个量对齐了:缓存条目的逻辑大小、它实际保活的 Storage,以及缓存预算用于驱逐的计数。对于文中的切分 View,复制放在“进入长期缓存”的位置,避免短期计算过程提前承担不必要的拷贝。

检查类似代码时,可以沿引用关系走一遍:谁创建了大 Tensor,哪些 View 共享它,哪个容器延长了引用生命周期,驱逐一个条目后还有谁保活底层存储。只有这些关系明确,才知道按 payload 计数、按 Storage 去重计数,还是在入库时分离存储更合适。

文中的连续 Batch 切分采用 strided Tensor。推广到其他缓存时,再根据 Entry 的共享关系与生命周期选择预算方式:有意复用的 Storage 可以按身份去重记账,需要独立驻留的数据则在入库时复制。

返回阅读地图 →
读到这里

从一条请求,看见整个系统。

返回文章目录

本页目录

搜索文章与 Skills

↑ ↓ 选择 · Enter 打开 · Esc 关闭 · Ctrl / ⌘ K 搜索

工程图