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

Clark11
18 阅读7分钟

目录

项目地址：

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 两个具体实现

实现类文件职责
MemBufferAllocatorCPUsrc/core/device/cpu/MemBufferAllocatorCPU.cppmalloc/free、memcpy、memset
MemBufferAllocatorCUDAsrc/core/device/cuda/MemBufferAllocatorCUDA.cppcudaMalloc/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 当前实现对齐。

avatar
文章
阅读
粉丝