GPU kernel 已返回，显存为什么还不能释放？用 Rust 管理 CPU/GPU 完整生命周期

在 GPU 推理代码里，最难维护的往往不是矩阵乘法，而是一个更朴素的问题：这块 内存到底属于谁，它什么时候可以释放？

Rust 的借用、Arc、Mutex 和 Drop 能把这个问题变得非常直接，但前提是先 分清三件不同的事：谁拥有资源、谁可以修改资源、GPU 什么时候使用完资源 。 如果不做区分，给所有对象统一套上 Arc<Mutex<_>>，代码虽然能够通过编译，却 会引入不必要的锁，甚至仍然可能在异步 kernel 尚未完成时提前释放显存。

zLLM 要建立一条贯穿所有 backend 的完整所有权链：

模型/文件中的数据
    │  Arc<File> / Vec<u8> / Vec<f32>
    ▼
RocmWeight 或 RocmTensor
    │  Arc<DeviceBuffer>
    ▼
DeviceBuffer（唯一显存 allocation owner，Drop 负责回收或释放）
    │  backend mempool 决定复用策略；Arc clone 被 in-flight 状态保留
    ▼
DeviceCompletion + HIP event（GPU 真正完成后才退休资源）

这套机制的关键不是"到处使用智能指针"，而是让每一种类型准确表达一种生命周期。 本文说明这套机制的定义、用法、优势和边界。

1. 先回答结论：不是 Arc<Mutex<Tensor>>

对于只读权重和大多数不可变 Tensor，推荐形态是：

Arc<DeviceBuffer>
Arc<RocmCooperativeMlaWeights> // 只有确实需要独立共享整个权重组时

而不是：

Arc<Mutex<DeviceBuffer>>
Arc<Mutex<RocmCooperativeMlaWeights>>

Arc 回答"有多少个使用者共同拥有它"，Mutex 回答"多个线程如何独占修改它"。 权重加载后通常只读，克隆 Arc 就能共享底层 allocation，没有必要为了读取指针 而加锁。

&RocmCooperativeMlaWeights 也不是资源 owner。它只表示当前函数临时借用一个 已经存在的权重组：

fn run_attention(weights: &RocmCooperativeMlaWeights) {
    // 函数调用期间可以读取权重；不取得整个权重组的所有权。
}

zLLM 把 cooperative MLA 权重保存在 RocmPrefillExperts 的 HashMap<usize, RocmCooperativeMlaWeights> 中，查询时返回引用：

pub(super) fn cooperative_mla_layer(
    &self,
    layer: usize,
) -> Result<(RocmContext, &RocmCooperativeMlaWeights), BackendError> {
    let weights = self.cooperative_mla.get(&layer)
        .ok_or_else(/* ... */)?;
    Ok((peer.context, weights))
}

这里的顶层引用已经足够，因为 RocmPrefillExperts 拥有权重组，调用者只在当前 计算中读取它。权重组内部的每个 RocmWeight 又通过 Arc<DeviceBuffer> 拥有 真正的显存。因此即使 RocmWeight 被 Clone，显存也不会重复上传或复制。

只有当整个权重组要脱离 RocmPrefillExperts，进入后台任务、队列或多个独立执行器 长期共享时，才值得把 map 的值升级为 Arc<RocmCooperativeMlaWeights>。

2. 第一层：CPU 内存和文件句柄的所有权

所有权链从模型文件开始，并遵循三种典型模式。

2.1 独占字节：Vec<u8> / Vec<f32>

Safetensors 的 TensorData 直接拥有读取出来的字节：

#[derive(Clone, Debug)]
pub struct TensorData {
    pub name: String,
    pub dtype: String,
    pub shape: Vec<usize>,
    pub data: Vec<u8>,
}

Vec 是最直接的 RAII owner：结构存在，字节存在；结构被释放，字节也被释放。 需要转换时，to_f32() 返回一个新的 Vec<f32>，所有权随返回值移动给调用者。 这种数据没有共享需求，不需要 Arc，更不需要 Mutex。

2.2 共享文件：Arc<File>

GGUF matrix 不会为每个 tensor 重新打开模型文件，而是保存共享文件句柄：

pub struct GgufMatrix {
    pub name: String,
    pub rows: usize,
    pub columns: usize,
    file: Arc<File>,
    offset: u64,
    bytes_len: usize,
    bytes: Arc<OnceLock<Vec<u8>>>,
}

矩阵切片通过 Arc::clone(&self.file) 共享文件，同时拥有独立的 offset 和 shape。 最后一个 matrix 释放后，文件句柄自然关闭。OnceLock<Vec<u8>> 则表达"原始字节 最多惰性加载一次"；读路径无需每次竞争互斥锁。

2.3 共享可变索引：Arc<Mutex<HashMap<...>>>

Safetensors shard cache 是 Mutex 真正适合出现的地方：

type SharedShards = Arc<Mutex<HashMap<String, Arc<CachedShard>>>>;

这里有一个会被多个 store 共同修改的 HashMap，所以既需要 Arc 共享目录，也 需要 Mutex 串行化插入。map 中的 shard 自身是 Arc<CachedShard>；拿到 shard 以后即可释放 map 锁，后续读取不必一直持锁。

全局缓存只保存 Weak，防止全局表永久拥有所有模型：

type SharedShardCaches =
    HashMap<PathBuf, Weak<Mutex<HashMap<String, Arc<CachedShard>>>>>;

这是一种值得统一采用的缓存模式：全局 registry 用 Weak 定位，真实使用者用 Arc 拥有，Mutex 只保护 registry 的结构变化。

3. 第二层：DeviceBuffer 是显存的唯一 owner

ROCm 路径中最重要的资源类型是 DeviceBuffer：

pub struct DeviceBuffer {
    device_id: i32,
    pointer: *mut c_void,
    bytes: usize,
    capacity_bytes: usize,
    recyclable: bool,
    async_allocated: bool,
    owner: Option<Arc<DeviceBuffer>>,
    // 省略 stage completion 与 deferred upload 状态
}

它同时记录设备、裸地址、逻辑大小、真实 allocation 容量、分配方式和可选上游 owner。HIP 裸指针只被关在这个底层类型里；backend 和 runtime 不直接负责调用 hipFree。

真正的释放策略集中在 Drop for DeviceBuffer：

impl Drop for DeviceBuffer {
    fn drop(&mut self) {
        if self.owner.is_some() {
            return; // view 不释放 owner 的 allocation
        }

        if self.async_allocated {
            // 优先在正确 stream 上 hipFreeAsync
        }

        if self.recyclable {
            // 能安全复用时归还显式 pool
        }

        // 其余情况最终 hipFree
    }
}

从上层看，分配和释放因此成为普通 Rust 值的创建与离开作用域：

let buffer = Arc::new(DeviceBuffer::allocate(device_id, bytes)?);
let another_user = Arc::clone(&buffer);

drop(buffer);       // allocation 仍存在
drop(another_user); // 最后一个 Arc 消失，才进入 DeviceBuffer::drop

RAII 在这里带来的不仅是"自动 free"，还把同步分配、异步分配、显式复用池和 view 统一到了同一个释放入口。调用者不需要记住这块地址最初来自 hipMalloc、 hipMallocAsync 还是 pool。

4. 第三层：Tensor 把逻辑信息与物理存储组合起来

ROCm Tensor 的定义是：

#[derive(Debug, Clone)]
pub struct RocmTensor {
    pub data: Vec<f32>,
    pub rows: usize,
    pub cols: usize,
    pub dtype: RocmTensorDType,
    pub layout: RocmTensorLayout,
    pub device: Option<Arc<DeviceBuffer>>,
}

它区分两类信息：

• rows、cols、dtype、layout 是逻辑 Tensor 描述；
• data 是必要时存在的 host shadow；
• device 是可选的设备 resident allocation。

构造一个纯设备 Tensor 时，上层只移动或共享 buffer：

fn device_tensor_with_arc(
    buffer: Arc<DeviceBuffer>,
    rows: usize,
    cols: usize,
    dtype: RocmTensorDType,
) -> RocmTensor {
    debug_assert_eq!(
        buffer.bytes(),
        rows * cols * dtype.element_bytes(),
    );
    RocmTensor {
        data: Vec::new(),
        rows,
        cols,
        dtype,
        layout: RocmTensorLayout::RowMajor,
        device: Some(buffer),
    }
}

RocmTensor::clone() 会复制少量 metadata、按需复制 host Vec，并克隆设备 buffer 的 Arc；它不会复制显存内容。只要任意 Tensor、权重、cache 或 in-flight 任务还持有这个 Arc，allocation 就不会进入 Drop。

这一层让 kernel 接口可以借用：

fn kernel(input: &RocmTensor, weight: &RocmWeight) -> Result<RocmTensor, Error>;

借用负责限制当前 CPU 调用栈中的访问；Tensor 内部的 Arc<DeviceBuffer> 负责底层 allocation 的共享所有权。两者职责不同，组合起来才完整。

5. 第四层：权重是不可变资源集合，不是互斥状态

RocmWeight 用枚举表达权重真正的 resident 形态。Dense 权重可以拥有 host F32 和设备副本；W4A16、W8A16、FP8、MXFP4、GGUF 则直接拥有 packed code、scale 等 Arc<DeviceBuffer>：

enum RocmWeightInner {
    Quantized(RocmQuantizedWeight),
    Dense {
        data: Vec<f32>,
        resident: Option<Arc<DeviceBuffer>>,
        resident_bf16: bool,
        router_bf16: Arc<OnceLock<Arc<DeviceBuffer>>>,
    },
    Gguf(GgufMatrix),
}

RocmCooperativeMlaWeights 再把同一执行能力所需的六个权重组合起来：

#[derive(Clone)]
struct RocmCooperativeMlaWeights {
    owner_q_b: RocmWeight,
    peer_q_b: RocmWeight,
    owner_kv_b: RocmWeight,
    peer_kv_b: RocmWeight,
    owner_o: RocmWeight,
    peer_o: RocmWeight,
}

这个 struct 的价值是表达"一层 cooperative MLA 的完整 resident 权重集合"，不是 提供并发控制。字段加载完成后保持不变，所以只读借用或 Arc 共享即可。

惰性生成的派生权重使用 OnceLock，例如 BF16 router cache：

Arc<OnceLock<Arc<DeviceBuffer>>>

这比 Mutex<Option<Arc<DeviceBuffer>>> 更精确：类型直接保证只初始化一次，初始化 完成后的热路径只读取结果。

6. View：子 Tensor 必须反向拥有原 allocation

GPU Tensor 经常只使用大 buffer 中的一段。如果 view 只保存偏移后的裸指针，原 buffer 一旦离开作用域就会释放，view 立即悬空。

zLLM 的 DeviceBuffer::view 把上游 owner 存进 view：

pub fn view(
    owner: Arc<DeviceBuffer>,
    offset: usize,
    bytes: usize,
) -> Result<DeviceBuffer, String> {
    let pointer = unsafe {
        owner.pointer.cast::<u8>().add(offset).cast()
    };
    Ok(DeviceBuffer {
        device_id: owner.device_id,
        pointer,
        bytes,
        owner: Some(owner),
        // view 自己不回收 allocation
        ..
    })
}

所有权关系变成：

RocmTensor
  └─ Arc<View DeviceBuffer>
       └─ Arc<Owner DeviceBuffer>
            └─ HIP allocation

view 的 Drop 发现 owner.is_some() 后不会释放裸地址；它只减少 owner 的引用 计数。最后一个 view 和原 owner 都消失后，真实 allocation 才会被释放或放回池。

这是用 Rust 封装显存后非常典型的收益：原本需要人工约定的"slice 不能活得比 base buffer 久"，现在由结构本身保证。

7. 最容易漏掉的一层：GPU 异步生命周期

到这里仍然没有完全解决问题。HIP kernel launch 通常是异步的：Rust 函数已经 返回、局部 Arc 已经 drop，并不代表 GPU 已经读完对应地址。

危险的伪代码如下：

fn launch(input: Arc<DeviceBuffer>) {
    unsafe { hip_launch(input.device_pointer()) };
} // input 在 CPU 侧释放，但 GPU 可能仍在执行

因此完整机制必须增加"完成所有权"。zLLM 用 DeviceCompletion 记录当前 stream 上的 HIP event，并持有这批异步工作仍依赖的资源：

pub struct DeviceCompletion {
    device_id: i32,
    event: usize,
    p2p_sources: Mutex<Vec<PendingP2pSource>>,
    stage_inputs: Vec<Arc<DeviceBuffer>>,
    recycled: Mutex<Vec<(usize, usize)>>,
    retired: AtomicBool,
}

提交阶段形成的临时输入与 P2P 来源被转移进 completion：

let completion = DeviceCompletion::record(device_id)?;

while !completion.is_complete()? {
    // CPU 可以处理别的请求
}
// event 完成后，p2p source、stage input 和待回收 allocation 才退休

对于 ordered P2P，代码会克隆来源 buffer 的 Arc，把它放进 PendingP2pSource，再由消费端 completion 持有。于是来源卡上的 allocation 至少 存活到复制或 peer read 真正完成，而不是只活到 hipMemcpyPeerAsync 返回。

Mutex 在 DeviceCompletion 中也是局部使用：它保护可能被轮询、等待或析构路径 共同退休的两个可变队列。它没有包住 Tensor，也没有包住 GPU 数据本身。

这给出一个非常重要的公式：

安全的 GPU 资源生命周期
  = Rust 所有权（Arc / Drop）
  + GPU 顺序（stream / event）
  + in-flight 引用保留（completion owns Arc）

缺少任何一项都不完整。Arc 不理解 GPU event，HIP event 也不会自动增加 Rust 引用计数；DeviceCompletion 是把两个世界接起来的桥。

8. 复用池：释放不等于马上 hipFree

高性能推理不能让每个临时 Tensor 都直接触发同步 hipFree。DeviceBuffer 的 Drop 会依据来源和状态选择：

1. view：只释放对 owner 的引用；
2. async allocation：优先 hipFreeAsync，保持 stream 顺序；
3. recyclable buffer：进入 stage 延迟回收或设备池；
4. 其他独占 allocation：最终调用 hipFree。

池本身用 Mutex<DeviceBufferPool> 保护 metadata：

struct DeviceBufferPool {
    buffers: HashMap<(i32, usize), Vec<usize>>,
    pending: HashMap<(i32, usize), Vec<PendingDeviceBuffer>>,
    bytes: HashMap<i32, usize>,
    available_events: Vec<usize>,
}

需要锁的是空闲地址列表、pending event 和计账，而不是 allocation 中的数十 MiB 数据。锁的作用域仅覆盖一次取出或归还，kernel 执行期间不会持锁。

这种设计把"逻辑释放"和"物理释放"分开：最后一个 Arc 消失代表业务不再需要 Tensor；Drop 再根据 event 与 pool 状态决定 allocation 是立即 free、异步 free 还是安全复用。上层算法无需知道这些差异。

9. Backend 无关的资源池契约

只分析 DeviceBufferPool 会产生一个错误印象：仿佛内存池是 ROCm/HIP 的特殊 优化。zLLM 把资源池定义为所有 backend 共同遵守的资源能力：

• ROCm 的 DeviceBufferPool 保存 available 与 pending allocation，并用 HIP event 判断何时可重新使用；
• Metal 的 MetalContext::scratch_pool 按 (tag, len, slot) 保存 metal::Buffer，利用同一 command queue 上的提交顺序复用 workspace；
• CPU Tensor 主要由 Vec<T> 拥有，普通 allocator 承担大部分工作，并不一定 需要 GPU 式 event pool。

所以 pool 的架构归属是 backend 资源管理能力 ，而不是 ROCm kernel 的公共 语义。backend/mod.rs 用统一的 MemoryPool trait、MemoryRequest、 MemoryKind 与 MemoryLifetime 定义契约。ROCm 的物理 pool 位于 kernel/rocm/hip/device_buffer.rs，Metal 的物理 pool 位于 backend/metal/context.rs；公共契约只统一请求和 RAII handle，不移动平台实现。

9.1 通用的应该是生命周期语义，不是裸指针实现

HIP、Metal、CUDA 和 CPU 的 allocation 不能被强行塞进一个最低公分母：

• HIP 有 device id、stream-ordered allocation、跨卡 event 和 P2P 生命周期；
• Metal 的 MTLBuffer 本身是引用计数对象，使用 command buffer/queue 表达完成；
• CUDA 有 stream、event、cudaMallocAsync 和设备 mempool；
• CPU 没有 GPU completion，alignment、NUMA、pinned memory 才是主要差异。

因此 backend 无关层不应持有 *mut c_void，也不应直接规定 hipFreeAsync 一类 平台动作。它真正需要统一的是五个问题：

申请：需要多少字节、对齐、用途和位置？
拥有：返回的 storage 由哪个 RAII handle 持有？
借出：同一 allocation 如何形成 Tensor/view？
退休：最后一个逻辑使用者退出后，设备是否已经完成？
复用：满足哪个 completion 条件后，allocation 可以重新进入 available？

统一接口这样表达这条边界：

trait MemoryPool {
    type Memory: Clone + Send + Sync + 'static;

    fn allocate_memory(
        &self,
        request: MemoryRequest,
    ) -> Result<Self::Memory, BackendError>;

    fn memory_bytes(&self, memory: &Self::Memory) -> u64;
}

struct MemoryRequest {
    bytes: usize,
    alignment: usize,
    kind: MemoryKind,
    lifetime: MemoryLifetime,
    tag: Option<&'static str>,
    slot: usize,
}

enum MemoryKind {
    Scratch,
    Activation,
    Cache,
    Weight,
    Transfer,
}

enum MemoryLifetime {
    Operation,
    Stage,
    Session,
    Model,
}

Metal 的 cached_zero_buffer_slot() 通过这一 trait 申请 Scratch；ROCm 的 concat_token_rows() 和 reserved concat allocation 通过同一 trait 申请 Operation/Stage Activation。接口刻意不增加统一 view() 或 retire_after()：view 由平台 storage 保留 owner，retirement 接到 StageExecutionBackend::Completion 和 backend RAII Drop，避免复制一套完成协议。

9.2 Pool 与 Arc 各管一半

Arc 与 mempool 不能相互替代：

Arc strong_count > 0
    → 仍有逻辑 owner，不能退休

Arc strong_count == 0
    → 进入 Storage::drop / retire
    → completion 未完成：进入 pending
    → completion 已完成：进入 available

下一次同类 allocation
    → 从 available best-fit/tag-slot 复用
    → 没有合适块才向设备申请新 allocation

Arc 决定"业务是否仍然引用资源"，pool 决定"业务释放后物理 allocation 去哪里"。 GPU completion 位于两者之间，决定资源何时从 pending 变成 available。这三者组合 才是一套完整的 backend 内存管理。

9.3 ROCm 与 Metal 的策略为什么不同

ROCm DeviceBufferPool 按 (device_id, capacity) 管理裸 allocation，并区分 buffers 与 pending。这是因为 host 端 drop 时 GPU 可能尚未执行完成，必须用 event 推进 pending；同时 best-fit 容量复用能适应不同 prefill shape。

Metal scratch_pool 按 (tag, len, slot) 管理 Buffer。tag 防止语义上可能同时 活跃的 workspace 意外别名，slot 则允许多个同尺寸切片同时存在。同一 queue 上 "写入、消费、下一次写入"的顺序为复用提供安全边界。

因此，backend 无关不等于所有 backend 使用相同 key 和相同容器：

层次	应统一	不应统一
runtime	Tensor 用途、生命周期等级、stage/session 边界	HIP event、Metal command buffer
backend capability	allocate/view/retire/completion 的行为契约	pool 的 key、best-fit 算法、平台句柄
backend 实现	本平台内一致的 owner 与回收路径	另一平台的同步原语
kernel	借用输入、写入输出	决定全局资源驻留和 pool 策略

9.4 权重、activation、KV 与 scratch 不能共用一种回收策略

统一 pool 还必须保留数据用途，而不能只看字节数：

• resident 权重：模型生命周期内常驻，通常不进入短期复用池；
• KV/DSA cache：属于 session，可增长、持久化或迁移，不能被 stage 结束回收；
• activation：跨若干 kernel 或 stage，由 completion 决定退休；
• scratch/workspace：生命周期最短，最适合按 tag/shape 高频复用；
• pinned host staging ：受 DMA completion 约束，不能按普通 Vec<u8> 处理。

也就是说，backend 无关的 mempool 首先需要统一"生命周期等级"，然后让 ROCm、 Metal、CUDA 或其他 backend 选择物理实现。若只提供 alloc(bytes) / free(ptr)， 上层仍然无法表达权重、KV 和临时 workspace 的根本区别。

9.5 统一架构边界

zLLM 的 mempool 遵守以下不变量：

1. 保留 RocmTensor、MetalTensor 等 backend associated type，不制造统一 dyn Tensor；
2. backend 生命周期边界统一描述 allocation 用途、stage completion 和 session retirement；
3. ROCm 用 event 驱动 pending → available，Metal 利用 command queue 顺序与 retained Buffer；
4. Tensor/storage handle 持有 pool return token，或在 Drop 中回到所属 pool；
5. 公共 capability 只表达真实共性，不建立通用裸指针 allocator，不让平台差异 反向污染 runtime。

完整关系是：

Backend
  ├─ Tensor / Weight / Cache associated types
  ├─ Completion（平台自己的 event/command buffer/fence）
  └─ Memory lifecycle policy
       ├─ allocate
       ├─ retain in flight
       ├─ retire to pending
       └─ promote/reuse/release

runtime 只表达资源用途与生命周期，不知道底层使用 HIP pool、Metal buffer cache、CUDA mempool 还是普通 CPU allocator。

10. 一套统一的使用规则

所有 backend 代码统一遵守下面的规则。

场景	推荐类型	原因
函数内临时读取 Tensor/权重	&RocmTensor、&RocmWeight	不转移所有权，作用域最短
多个 Tensor/权重共享一块显存	Arc<DeviceBuffer>	最后一个使用者负责触发回收
共享整组只读权重	Arc<Weights>，或由单一 manager 拥有并返回 &Weights	不需要锁
只初始化一次的派生权重	OnceLock<Arc<DeviceBuffer>>	语义比互斥 Option 更准确
可变 cache/registry metadata	Mutex<HashMap<...>>	只保护结构变化
全局但不应永久保活的 cache	Weak<T> registry + Arc<T> user	自动随最后一个用户退出
大 buffer 的逻辑切片	view 内保存 Arc<owner>	防止 base allocation 提前释放
异步 kernel/P2P 依赖	completion 持有 Arc<DeviceBuffer>	直到 event 完成才退休
独占、短生命周期 CPU 数据	Vec<T> / Box<T>	避免无意义引用计数

新增 GPU 接口时，可以按下面的顺序审查：

1. 裸 pointer 是否只存在于 DeviceBuffer 或 kernel FFI 边界？
2. 返回 view 时是否保留了 owner？
3. Arc 的最后一个 CPU 引用消失时，GPU 是否可能仍在使用它？
4. 如果可能，哪个 completion/event 持有它？
5. Mutex 是否只保护确实可变的 metadata，锁内是否包含 I/O、同步或 kernel 等待？
6. cache 是否需要 Weak 防止永久驻留？

11. 一个从加载到异步执行的完整例子

下面的简化例子展示这套机制如何连起来：

// 1. 上传后，Arc 成为 resident allocation 的共享 owner。
let q_b = Arc::new(DeviceBuffer::upload(device_id, q_b_bytes)?);

// 2. 权重只读，不加 Mutex。
let weights = Arc::new(MlaWeights {
    q_b: Arc::clone(&q_b),
});

// 3. Tensor view 反向持有 q_b，避免基址先释放。
let shard = Arc::new(DeviceBuffer::view(
    Arc::clone(&weights.q_b),
    shard_offset,
    shard_bytes,
)?);

// 4. kernel 只借用资源。
launch_mla(input_buffer.as_ref(), shard.as_ref())?;

// 5. 如果 GPU 工作越过当前 Rust 作用域，submission 必须取得 Arc。
submission.retain(Arc::clone(&input_buffer));
submission.retain(Arc::clone(&shard));
let completion = DeviceCompletion::record(device_id)?;

// 6. 业务 owner 可以先退出；in-flight owner 仍保活 allocation。
drop(weights);
drop(shard);

// 7. event 完成后，completion 释放引用；最后一个 Arc 触发 Drop，
//    DeviceBuffer 决定归还 pool、hipFreeAsync 或 hipFree。
completion.wait()?;

这里的 submission.retain() 表达统一语义；各 backend 把它落实到自己的 in-flight 资源集合。ROCm 由 active stage、ordered P2P pending source 和 DeviceCompletion 完成保活，Metal 由 command buffer 的 retained resources 完成保活。

12. 这套机制带来的实际优势

生命周期从约定变成类型关系

RocmTensor → Arc<DeviceBuffer>、View → Arc<Owner>、 Completion → Vec<Arc<DeviceBuffer>> 都可以直接从字段看出来。开发者不再依赖 "调用者应该记得晚一点 free"这样的口头约定。

显著减少重复上传和复制

克隆 RocmWeight 或 RocmTensor 只增加设备 buffer 的引用计数，权重显存不会 被复制。量化 codes、scales 和派生 BF16 cache 也能独立共享。

错误路径自动收敛

函数在中途使用 ? 返回时，已经构造成功的 Vec、Arc、Tensor 和 buffer 会按 逆序析构。大部分错误分支不再需要手写多段 cleanup。

资源策略在 backend 内独立实现

runtime 只使用 backend 的 Tensor/Weight/Completion 契约；ROCm 使用 hipFreeAsync、显式 pool、stage 批量回收和 event 驱动退休，Metal 使用 command queue 与 scratch pool，两者都不改变模型算法。

并发边界更清楚

不可变数据用 Arc，一次性初始化用 OnceLock，可变 metadata 才用 Mutex。 锁竞争因此集中在小型 registry 和队列，而不是扩散到权重读取热路径。

13. 仍然需要保持警惕的边界

这套机制很强，但不是自动正确。

• unsafe impl Send/Sync for DeviceBuffer 是对底层 HIP 行为的人工承诺，编译器无法 替我们证明跨线程和跨设备操作安全；每条新 FFI 路径仍需遵守 device/stream 规则。
• Arc 只能防止 Rust 对象析构，不能自动建立 producer/consumer stream 顺序。 数据依赖仍要用 stream ordering 或 event 表达。
• host shadow 与 device resident 同时存在时，要明确谁是最新副本；引用计数不能 解决数据一致性。
• Arc 环会导致资源永久不释放。owner 方向应保持单向：view 指向 allocation， allocation 不反向拥有 view；registry 需要时使用 Weak。
• Drop 不适合返回错误。真正要求"确认 GPU 已完成"的边界，应显式调用 wait() 或由 scheduler 轮询 completion，析构只作为最后的安全网。
• 不应为了形式统一把所有 Tensor 都改成 Arc<Tensor>。通常共享的是大块 storage， Tensor metadata 本身很小，按值移动或克隆更简单。

结语

Rust 可以极大降低 CPU 内存与 GPU 显存的生命周期管理复杂度。zLLM 的统一答案是：

Vec / Arc<File> 管理 host 数据与文件
RocmWeight / RocmTensor 描述逻辑资源
Arc<DeviceBuffer> 共享物理显存
Drop 统一回收、异步释放与复用策略
owner Arc 保证 view 不悬空
DeviceCompletion + HIP event 保证 GPU 用完以后才退休
Mutex 只保护真正可变的 registry 与回收队列
backend mempool 统一生命周期语义，各平台保留自己的分配与完成原语

因此，真正值得推广的不是 Arc<Mutex<Tensor>> 这个单一写法，而是这套分层的 所有权协议：借用表达当前调用，Arc 表达共享存活，Mutex 表达可变互斥， Drop 表达资源归还，event/completion 表达设备完成。 每个原语只做自己擅长的 一件事，组合起来才能同时得到安全、性能和较低的开发复杂度。

---

原文：zhuai.tech/blog/rust-u...

zLLM GitHub：github.com/zllm-lab/zl...