手写LLM推理框架时，内存管理99%的人会踩的坑 | TFFInfer解析(五)——Tensor 张量系统与内存抽象（下）

目录

• [1. Memory 类：raw pointer 的 RAII 封装](#1. Memory 类：raw pointer 的 RAII 封装 "#1-memory-%E7%B1%BBraw-pointer-%E7%9A%84-raii-%E5%B0%81%E8%A3%85")
• [2. 分配器体系：从抽象到实现](#2. 分配器体系：从抽象到实现 "#2-%E5%88%86%E9%85%8D%E5%99%A8%E4%BD%93%E7%B3%BB%E4%BB%8E%E6%8A%BD%E8%B1%A1%E5%88%B0%E5%AE%9E%E7%8E%B0")
• [3. CPU 与 GPU 分配器的差异](#3. CPU 与 GPU 分配器的差异 "#3-cpu-%E4%B8%8E-gpu-%E5%88%86%E9%85%8D%E5%99%A8%E7%9A%84%E5%B7%AE%E5%BC%82")
• [4. MemManager 与 Tensor 的协作关系](#4. MemManager 与 Tensor 的协作关系 "#4-memmanager-%E4%B8%8E-tensor-%E7%9A%84%E5%8D%8F%E4%BD%9C%E5%85%B3%E7%B3%BB")
• [5. 本节小结](#5. 本节小结 "#5-%E6%9C%AC%E8%8A%82%E5%B0%8F%E7%BB%93")

---

项目地址：

GitHub：github.com/NKKdev/TFFi...

Gitee：gitee.com/NKK_Ovit/tf...

在自研的3万行LLM推理框架TFFInfer中，内存管理是最大的性能陷阱。一个错误的分配策略，可能导致显存崩溃或推理延迟暴增10倍。这篇文章，我会手把手拆解如何用RAII和分配器多态，写出工业级的内存管理模块。

1. Memory 类：raw pointer 的 RAII 封装

上一节我们讲了 Tensor 如何描述「数据长什么样」，这一节讲它背后的「数据存在哪里」------Memory 类。

1.1 设计定位

Memory（src/core/mem/Memory.h）是一个轻量级的 RAII 内存句柄，职责非常纯粹：

• 持有指向实际数据的 void*
• 记录字节大小
• 在析构时通过分配器释放内存
• 支持外部指针（不释放）和内部指针（自动释放）两种模式

1.2 核心实现

class Memory : public std::enable_shared_from_this<Memory> {
    size_t _byte_size = 0;
    void* _ptr = nullptr;
    bool _use_external = false;        // 是否外部指针（不由我释放）
    DeviceType _device_type = TFF_BACKEND_DEVICE_TYPE_UNKNOWN;
    std::shared_ptr<MemBufferAllocatorBaseObject> _allocator;
    bool _is_used = false;             // 占用标记（用于内存池管理）
};

构造函数：

explicit Memory(size_t byte_size, void* ptr = nullptr, bool use_external = false,
                std::shared_ptr<device::MemBufferAllocatorBaseObject> allocator = nullptr) {
    this->_byte_size = byte_size;
    this->_allocator = std::move(allocator);
    if (use_external) {
        this->_ptr = ptr;              // 直接接管外部指针
        this->_use_external = true;
        this->_is_used = true;
    }
    this->reset();
}

析构函数------RAII 的核心：

virtual ~Memory() {
    if (!_use_external && _allocator != nullptr) {
        _allocator->release(_ptr);     // 通过分配器释放
        this->_byte_size = 0;
        _allocator = nullptr;
    }
}

关键洞察 ：Memory 本身不直接调用 free 或 cudaFree，而是委托给 _allocator。这使得同一块 Memory 对象可以在不知道自己是 CPU 内存还是 GPU 显存的情况下正确释放------完全的多态行为。

1.3 allocate() 与 copy_from()

bool Memory::allocate() {
    if (_allocator && _byte_size > 0) {
        _ptr = _allocator->allocate(_byte_size);
        if (!_ptr) {
            _allocator->memset_zero(_ptr, _byte_size);
        }
    }
    return _ptr != nullptr;
}

copy_from() 的设计很有趣------它自动判断设备类型并选择拷贝方向：

void Memory::copy_from(const Memory &mem) {
    if (mem._allocator && mem._byte_size > 0) {
        this->_allocator = mem._allocator;
        this->_byte_size = mem._byte_size;
      
        if (this 是 GPU && mem 是 CPU) {
            mem._allocator->memcopy(mem._ptr, this->_ptr, this->_byte_size, 
                                    TFF_MEM_CPY_TYPE_HOST2DEVICE);
        } else if (this 是 CPU && mem 是 GPU) {
            mem._allocator->memcopy(mem._ptr, this->_ptr, this->_byte_size,
                                    TFF_MEM_CPY_TYPE_DEVICE2HOST);
        }
    }
}

注意：这里有个微妙的 bug 风险------this->_ptr 在拷贝前可能还没分配。实际使用中，通常是先 allocate() 再 copy_from()。这个设计体现了「让对象自己知道如何处理跨设备拷贝」的思想，上层无需关心 CUDA/HIP 细节。

1.4 is_used / reset / occupy

这三个方法服务于内存池的占用标记：

inline void reset()  { this->_is_used = false; }   // 标记为空闲
inline void occupy() { this->_is_used = true; }    // 标记为占用

在 MemManager 的内存复用逻辑中，一个张量生命周期结束后，其底层 Memory 会被 reset()，但不会立即释放物理内存 ，而是留在池中等待被下一个同尺寸需求 occupy()。

---

2. 分配器体系：从抽象到实现

2.1 基类设计

MemBufferAllocatorBaseObject（src/core/device/MemBufferAllocatorBaseObject.h）定义了所有分配器必须实现的接口：

class MemBufferAllocatorBaseObject : public ModuleObject {
public:
    virtual void* allocate(size_t size) = 0;           // 分配
    virtual void release(void* ptr) = 0;               // 释放
    virtual void memset_zero(void* ptr, size_t size) = 0;  // 清零
    virtual void memcopy(void* src, void* dst, size_t size, MemCpyKind kind) = 0;
  
    DeviceType _device_type = TFF_BACKEND_DEVICE_TYPE_UNKNOWN;
    int _device_id = -1;
};

继承自 ModuleObject 意味着它可以通过 ModuleFactory 注册和创建------TFFInfer 中几乎所有核心组件都走工厂路线。

2.2 两个具体实现

实现类	文件	职责
MemBufferAllocatorCPU	src/core/device/cpu/MemBufferAllocatorCPU.cpp	malloc/free、memcpy、memset
MemBufferAllocatorCUDA	src/core/device/cuda/MemBufferAllocatorCUDA.cpp	cudaMalloc/cudaFree、cudaMemcpy、cudaMemset

CPU 分配器：

void* MemBufferAllocatorCPU::allocate(size_t size) {
    void* ptr = std::aligned_alloc(ALIGNMENT, size);  // 对齐分配
    return ptr;
}

void MemBufferAllocatorCPU::release(void* ptr) {
    std::free(ptr);
}

void MemBufferAllocatorCPU::memcopy(void* src, void* dst, size_t size, MemCpyKind kind) {
    std::memcpy(dst, src, size);  // CPU 间拷贝，kind 无意义
}

CUDA 分配器：

void* MemBufferAllocatorCUDA::allocate(size_t size) {
    void* ptr = nullptr;
    cudaMalloc(&ptr, size);       // 设备显存分配
    return ptr;
}

void MemBufferAllocatorCUDA::release(void* ptr) {
    cudaFree(ptr);
}

void MemBufferAllocatorCUDA::memcopy(void* src, void* dst, size_t size, MemCpyKind kind) {
    cudaMemcpy(dst, src, size, convert_cuda_memcpy_kind(kind));
}

2.3 对齐的重要性

两个分配器都使用 ALIGNMENT（通常为 256 字节）对齐：

• GPU 全局内存合并访问：未对齐的地址可能导致 warp 内线程访问不同 cache line，带宽利用率暴跌。
• DMA 传输：某些 CUDA copy engine 要求源/目的地址对齐。
• 量化 kernel ：Q8_0 的 block 读取常假设地址对齐。

---

3. CPU 与 GPU 分配器的差异

3.1 延迟差异

操作	CPU (malloc)	GPU (cudaMalloc)
分配延迟	~μs 级	~ms 级（驱动与硬件交互）
释放延迟	~μs 级	~ms 级
频繁分配影响	较小	极大（导致 GPU 卡顿）

这就是 TFFInfer 设计显存池 的根本原因------绝不在推理主路径上调用 cudaMalloc/cudaFree。

3.2 内存拷贝种类

enum class MemCpyKind {
    TFF_MEM_CPY_TYPE_HOST2HOST,    // CPU → CPU
    TFF_MEM_CPY_TYPE_HOST2DEVICE,  // CPU → GPU
    TFF_MEM_CPY_TYPE_DEVICE2HOST,  // GPU → CPU
    TFF_MEM_CPY_TYPE_DEVICE2DEVICE,// GPU → GPU
};

MemBufferAllocatorCUDA::memcopy 内部映射到 CUDA 的四种 cudaMemcpyKind：

• cudaMemcpyHostToHost
• cudaMemcpyHostToDevice
• cudaMemcpyDeviceToHost
• cudaMemcpyDeviceToDevice

3.3 页锁定内存（Pinned Memory）

当前 TFFInfer 的 CUDA 分配器使用普通 cudaMalloc。后续若优化 CPU→GPU 的权重加载速度，可以引入 cudaMallocHost（页锁定内存）：

• 页锁定内存 ↔ 设备内存 的拷贝速度比 可分页内存 快 2~3 倍
• 代价：页锁定内存无法被操作系统换出，过量使用会降低系统整体内存性能

这是一个典型的优化取舍点------TFFInfer 当前选择了简单性，保留了升级空间。

---

4. MemManager 与 Tensor 的协作关系

4.1 谁负责分配？

这是一个容易混淆的问题。在 TFFInfer 中：

层级	职责
Tensor	描述「我要多大的内存」，调用allocate() 或 set_buffer_data()
Memory	持有指针，在析构时释放
MemBufferAllocator*	实际的malloc/cudaMalloc 调用者
LLMMemManager	显存池管理者，决定「从哪里分配、什么时候回收、怎么复用」

典型流程：

// 1. Tensor 知道自己需要多少字节
auto tensor = std::make_shared<Tensor>(DataType::TFF_DATA_TYPE_F32, 
                                       MemoryType::TFF_MEM_TYPE_WORKSPACE,
                                       shapes, true, nullptr);

// 2. MemManager 从池中找一块合适的空闲内存
auto [offset, ptr] = mem_manager->allocate_memory(tensor->get_bytes(), 
                                                   device_id, 
                                                   MemoryType::TFF_MEM_TYPE_WORKSPACE,
                                                   event);

// 3. Tensor 绑定这块外部内存
tensor->set_buffer_data(ptr, tensor->get_bytes(), offset);

4.2 分配器如何绑定到 Tensor

在 Tensor 的构造函数中：

// 内部分配模式
tensor.allocate();  
// → 创建 Memory，Memory 调用 _allocator->allocate()

// 外部绑定模式
tensor.set_buffer_data(ptr, size, offset);
// → 创建 Memory，传入外部 ptr，_use_external = true
// → 同时可以 set_allocator(allocator) 用于后续释放

4.3 一个容易忽略的细节

Tensor::release()：

inline void release() {
    if (_buffer) {
        _allocator->release(_buffer->ptr());
    }
    _type_size = 0;
    _blk_size = 0;
}

注意：这里直接调用了 _allocator->release()，而不是等 Memory 析构。这是因为显存池场景中，回收的是内存池中的 offset，而不是释放物理内存 。MemManager 重写了 release() 的行为------它只是把这块区域标记为空闲，而不是真的 cudaFree。

---

5. 小结

5.1 核心要点

1. Memory 是 RAII 句柄 ：封装 void* + 字节大小 + 分配器，析构时自动释放。
2. 分配器多态 ：MemBufferAllocatorCPU/CUDA 统一接口，Tensor/Memory 无需感知设备差异。
3. 对齐是性能基础：256 字节对齐服务于 GPU 内存合并访问和 DMA。
4. 显存池的必要性 ：cudaMalloc 延迟 ms 级，推理主路径必须通过池化避免。
5. MemManager 是「指挥官」：Tensor 申请、Memory 持有、Allocator 执行、MemManager 调度。

5.2 思考题

1. Memory::copy_from() 中如果 this->_ptr 为 nullptr 会发生什么？实际代码中应该如何防御？
2. 为什么 MemBufferAllocatorCUDA::allocate 不使用 cudaMallocManaged（统一内存）？统一内存的优缺点是什么？
3. 假设层 A 产出的张量生命周期为 [2, 5]，层 B 需要的张量生命周期为 [6, 8]，且两者大小相同。MemManager 如何将 A 的内存复用给 B？需要修改哪些数据结构？

5.3 预告

第 5 课 将跳出单个张量的视角，进入计算图层面：

• Graph 如何表示 DAG（有向无环图）
• GraphNode 的输入输出连接机制
• 拓扑排序与生命周期分析的完整流程
• 环检测的实现

---

文档版本：与仓库 src/core/mem/Memory.h、src/core/device/MemBufferAllocatorBaseObject.h 当前实现对齐。