# 十年 C++ 后端 GAP 六个月，写了一个近 3 万行的LLM-TFFInfer推理框架项目解析(四)——Tensor 张量系统与内存抽象（上）

# 十年 C++ 后端 GAP 六个月，写了一个近 3 万行的LLM-TFFInfer推理框架项目解析

# 十年 C++ 后端 GAP 六个月，写了一个近 3 万行的LLM-TFFInfer推理框架项目解析(二)

#十年 C++ 后端 GAP 六个月，写了一个近 3 万行的LLM-TFFInfer推理框架项目解析(三)-模型加载

知乎-Tffinfer框架源码解析

目录

• [1. Tensor](#1. Tensor "#1-%E4%B8%BA%E4%BB%80%E4%B9%88%E5%85%88%E8%AE%B2-tensor")
• [2. 类型系统：DataType、MemoryType、ModelTensorType](#2. 类型系统：DataType、MemoryType、ModelTensorType "#2-%E7%B1%BB%E5%9E%8B%E7%B3%BB%E7%BB%9Fdatatypememorytypemodeltensortype")
• [3. Tensor 类的整体设计](#3. Tensor 类的整体设计 "#3-tensor-%E7%B1%BB%E7%9A%84%E6%95%B4%E4%BD%93%E8%AE%BE%E8%AE%A1")
• [4. Shape、Stride 与多维索引](#4. Shape、Stride 与多维索引 "#4-shapestride-%E4%B8%8E%E5%A4%9A%E7%BB%B4%E7%B4%A2%E5%BC%95")
• [5. 张量生命周期与内存复用](#5. 张量生命周期与内存复用 "#5-%E5%BC%A0%E9%87%8F%E7%94%9F%E5%91%BD%E5%91%A8%E6%9C%9F%E4%B8%8E%E5%86%85%E5%AD%98%E5%A4%8D%E7%94%A8")
• [6. 本节小结与下节预告](#6. 本节小结与下节预告 "#6-%E6%9C%AC%E8%8A%82%E5%B0%8F%E7%BB%93%E4%B8%8E%E4%B8%8B%E8%8A%82%E9%A2%84%E5%91%8A")

---

1.Tensor

如果说模型加载 是推理框架的「入口」，那么 Tensor（张量） 就是整个框架的「血液」。

在 TFFInfer 中，从模型权重的存储、中间激活值的传递、到 KV Cache 的缓存，无一不是以 Tensor 为基本单元流转的。理解 Tensor 的设计，等于理解了框架中「数据长什么样、存在哪里、怎么被访问」这三个核心问题。

对于AI从业者来说，Tensor 这个概念应该不陌生------PyTorch/TensorFlow 里天天打交道。但手写一个工业级的 C++ Tensor 类，需要考虑的问题远比 Python 接口多得多：

• 如何支持 FP32/FP16/INT8/Q8_0 等多种数据类型？
• 如何表示量化张量（block size 怎么处理）？
• 如何计算多维索引偏移（Stride）？
• 如何区分「自己管理的内存」和「外部绑定的内存」？
• 如何为显存复用提供生命周期信息？

本节将带你逐层拆解 TFFInfer 中 Tensor 的设计哲学。

---

2. 类型系统：DataType、MemoryType、ModelTensorType

在深入 Tensor 类之前，我们先理清框架中的三个核心枚举类型。它们分别回答：数据是什么格式？内存是什么用途？张量在模型中扮演什么角色？

2.1 DataType ------ 数据类型

定义在 src/core/mem/BaseDefine.h 中：

enum class DataType {
    TFF_DATA_TYPE_UNKNOWN = 0,
    TFF_DATA_TYPE_F32,    // FP32
    TFF_DATA_TYPE_F16,    // FP16
    TFF_DATA_TYPE_Q8_0,   // Q8_0 量化（每 block 256 个元素配 1 个 scale）
    TFF_DATA_TYPE_Q4_0,   // Q4_0 量化（预留）
    TFF_DATA_TYPE_I32,    // INT32
    TFF_DATA_TYPE_I16,    // INT16
    TFF_DATA_TYPE_I8,     // INT8
};

每个 DataType 都对应一个 TypeTraits，描述该类型的元素大小 和块大小：

struct TypeTraits {
    size_t _type_size;   // 单个元素的字节数（如 FP32 = 4）
    size_t _blck_size;   // 量化 block 大小（非量化为 1）
};

为什么要设计 _blck_size？

以 Q8_0 为例，它不是每个元素独立存储 8bit，而是每 256 个 FP32 元素量化成 256 个 INT8 + 1 个 scale 浮点数。这意味着：

• 逻辑上 _shape[0] = 4096
• 物理上实际占用的字节数 = 4096 + 4096/256 * sizeof(float)
• Stride 计算时必须考虑 _blck_size

2.2 MemoryType ------ 内存用途类型

enum class MemoryType {
    TFF_MEM_TYPE_RESIDENT,   // 常驻内存（如权重、RoPE 表）
    TFF_MEM_TYPE_WORKSPACE,  // 工作区（中间激活值）
    TFF_MEM_TYPE_KV_CACHE,   // KV Cache 专用
};

这个分类直接影响内存分配策略：

• RESIDENT：预分配、永不释放（模型权重）
• WORKSPACE：按生命周期动态分配与回收
• KV_CACHE：Paged 管理，按 page 粒度分配

2.3 ModelTensorType ------ 模型语义类型

enum class ModelTensorType {
    LLM_TENSOR_TYPE_UNKNOWN,
    LLM_TENSOR_TOKEN_EMBD,      // 词嵌入权重
    LLM_TENSOR_ATTN_Q,          // Attention Q 权重
    LLM_TENSOR_ATTN_K,          // Attention K 权重
    LLM_TENSOR_ATTN_V,          // Attention V 权重
    LLM_TENSOR_ATTN_OUT,        // Attention Output 权重
    LLM_TENSOR_FFN_GATE,        // FFN Gate 权重
    LLM_TENSOR_FFN_UP,          // FFN Up 权重
    LLM_TENSOR_FFN_DOWN,        // FFN Down 权重
    LLM_TENSOR_ATTN_NORM,       // Attention 前 RMSNorm 权重
    LLM_TENSOR_FFN_NORM,        // FFN 前 RMSNorm 权重
    LLM_TENSOR_OUTPUT_NORM,     // 输出层 RMSNorm 权重
    LLM_TENSOR_OUTPUT,          // 输出 LM Head 权重
};

这个枚举在模型加载 阶段由 get_model_tensor_type() 根据张量名字解析得到，随后在图构建 阶段指导 ModelCreator 将权重绑定到正确的算子节点上。

---

3. Tensor 类的整体设计

Tensor 定义在 src/core/mem/Tensor.h，继承自 std::enable_shared_from_this<Tensor>。这个继承很关键------框架中大量通过 shared_ptr<Tensor> 传递张量，需要支持循环引用检测和安全的生命周期管理。

3.1 核心成员变量

class Tensor : public std::enable_shared_from_this<Tensor> {
    // === 维度信息 ===
    size_t _n_dims;                                    // 有效维度数（去尾1后）
    std::array<int64_t, MAX_TENSOR_DIM> _shape;        // 形状（最多4维）
    std::array<int64_t, MAX_TENSOR_DIM> _strides;      // 各维步长（字节）

    // === 类型信息 ===
    DataType _data_type;                               // 数据类型
    size_t _type_size;                                 // 元素大小（来自 TypeTraits）
    uint32_t _blk_size;                                // 块大小（量化用）

    // === 内存管理 ===
    MemoryType _memory_type;                           // 内存用途类型
    std::shared_ptr<Memory> _buffer;                   // 底层内存对象
    bool _use_external;                                // true：外部内存；false：自己分配
    int64_t _external_memory_index;                    // 外部内存池中的索引
    std::shared_ptr<MemBufferAllocatorBaseObject> _allocator;  // 分配器

    // === 生命周期（用于显存复用） ===
    int _start;                                        // 活跃开始时间点
    int _end;                                          // 活跃结束时间点
    int _priority;                                     // 优先级（分配时参考）

    // === 模型语义 ===
    ModelTensorType _tensor_type;                      // 模型中的角色
};

3.2 两种构造模式

Tensor 支持两种内存管理模式，这是工业级框架的必备设计：

模式 A：内部自主分配（use_external = false）

Tensor(DataType::TFF_DATA_TYPE_F32, 
       MemoryType::TFF_MEM_TYPE_WORKSPACE,
       std::array<int64_t, 4>{4096, 1, 1, 1},
       false,                       // 自己分配
       allocator);
// 构造函数内部自动调用 allocate()

模式 B：外部内存绑定（use_external = true）

Tensor(DataType::TFF_DATA_TYPE_F32,
       MemoryType::TFF_MEM_TYPE_RESIDENT,
       std::array<int64_t, 4>{4096, 11008, 1, 1},
       true,                        // 外部内存
       nullptr);

// 后续由 MemManager 统一分配大 buffer 后绑定
tensor.set_buffer_data(gpu_ptr, byte_size, memory_offset);

为什么要区分这两种模式？

场景	模式	原因
模型权重	外部	权重总量大，需要统一规划显存布局，且生命周期贯穿始终
中间激活值	外部	由MemManager 按生命周期复用同一块显存
单元测试/临时张量	内部	简单场景，无需复杂管理

---

4. Shape、Stride 与多维索引

4.1 Stride 推导算法

这是手写 Tensor 最核心的算法之一。给定 _shape 和 _type_size（及 _blk_size），计算每个维度的步长：

inline void stride_infer() {
    // 第0维：元素大小（考虑量化块）
    _strides[0] = _type_size;
  
    // 第1维：第0维的总字节跨度
    if (_strides.size() > 1) {
        _strides[1] = _strides[0] * (_shape[0] / _blk_size);
    }
  
    // 第2维及以上：累积乘积
    for (int j = 2; j < _shape.size(); ++j) {
        _strides[j] = _strides[j - 1] * _shape[j - 1];
    }
}

举例 ：一个形状为 [4096, 11008] 的 FP32 权重矩阵

• _strides[0] = 4（FP32 元素大小）
• _strides[1] = 4 * 4096 = 16384
• 元素 (i, j) 的偏移 = i * 4 + j * 16384

量化举例 ：Q8_0 类型，_blk_size = 256，形状 [4096, 11008]

• 逻辑上仍是 4096 × 11008 个元素
• 但 _strides[1] = 4 * (4096 / 256) = 64（因为每 256 个元素压缩成一个 block）
• 实际内存布局由 CUDA kernel 的量化格式决定，Stride 只服务于索引计算

4.2 类型安全的元素访问

template<typename T, typename... Args>
inline T &at(Args... indices) {
    if (sizeof(T) != _type_size) {
        throw std::runtime_error("Tensor::get<T>(): sizeof(T) != element size");
    }
    if (!_buffer || !_buffer->ptr()) {
        throw std::runtime_error("Tensor buffer is null");
    }
    return *reinterpret_cast<T*>(_buffer->ptr() + compute_offset(indices...));
}

compute_offset 将多维索引展开为一维字节偏移：

template<typename... Args>
inline size_t compute_offset(Args... indices) const {
    constexpr size_t num_indices = sizeof...(indices);
    std::array<size_t, num_indices> idxs{static_cast<size_t>(indices)...};
    size_t offset = 0;
    for (size_t i = 0; i < num_indices; ++i) {
        offset += idxs[i] * _strides[i];
    }
    return offset;
}

注意 ：at() 虽然方便，但在推理主路径上不会被调用------CUDA kernel 直接操作 raw pointer。它主要用于：

• 调试时打印张量内容
• CPU 侧的后处理（如采样前的 logits 读取）
• 单元测试验证

4.3 check()：自动压缩尾维

inline void check() {
    this->_n_dims = _shape.size();
    for (int i = _shape.size() - 1; i >= 0; --i) {
        if (_shape[i] == 1) {
            this->_n_dims--;
        } else {
            break;
        }
    }
}

这个函数会在构造时自动去掉尾部为 1 的维度。例如传入 shape [4096, 1, 1, 1]，实际 _n_dims = 1。这减少了后续算子中不必要的维度处理。

---

5. 张量生命周期与内存复用

5.1 生命周期标记

inline void set_live_range(int start, int end) {
    this->_start = start;
    this->_end = end;
}

_start 和 _end 是执行时间戳 ，由 Graph::analyze_lifetimes() 在拓扑排序后计算：

void Graph::analyze_lifetimes() {
    // 为每个节点分配执行时间戳
    for (size_t i = 0; i < _nodes.size(); ++i) {
        _exec_time[_nodes[i]] = i;
    }
  
    // 计算每个张量的首次和最后使用时间
    for (auto& node : _nodes) {
        for (auto& output : node->outputs()) {
            int exec_time = _exec_time[node];
            if (_first_use.find(output) == _first_use.end()) {
                _first_use[output] = exec_time;   // 首次产生
            }
            _last_use[output] = exec_time;        // 最后被使用
        }
    }
}

5.2 生命周期如何驱动显存复用

这是 TFFInfer 显存优化的核心。假设一个 Transformer 层的计算流程：

时间戳	算子	产生的张量	生命周期
0	Embedding	emb_out	[0, 1]
1	RMSNorm	norm_out	[1, 2]
2	MatMul(Q)	q_out	[2, 5]
3	MatMul(K)	k_out	[3, 5]
4	MatMul(V)	v_out	[4, 5]
5	FlashAttention	attn_out	[5, 6]
6	Add(Residual)	add_out	[6, ...]

观察发现：

• emb_out 在 t=1 后就不再被需要
• norm_out 在 t=2 后就不再被需要
• q_out、k_out、v_out 在 t=5 后就不再被需要

MemManager 会根据这些生命周期，把已经「死亡」的张量内存标记为可回收，供后续算子复用。这样 28 层的 Transformer，不需要为每层单独分配激活显存，而是可以复用同一块物理内存。

---

6. 本节小结与下节预告

6.1 核心要点回顾

1. 类型系统三层 ：DataType 管格式、MemoryType 管用途、ModelTensorType 管语义，三者正交组合描述一个张量。
2. 两种内存模式：内部自主分配适合简单场景，外部绑定适合大规模显存池管理。
3. Stride 推导 ：量化场景下需考虑 _blk_size，是手写 Tensor 最容易出错的细节。
4. 生命周期标记 ：_start/_end 是显存复用的元数据基础，由 Graph 在编译期推导。

6.2 思考

1. 若一个 Q8_0 张量的 _blk_size = 256，_shape = {512, 1024}，_strides[1] 应该等于多少？实际物理字节数又是多少？
2. 为什么 Tensor 继承 std::enable_shared_from_this？如果不继承，在 get_ref_count() 中可能遇到什么异常？
3. check() 函数压缩尾维后，对一个 shape 为 [1, 1, 128, 1] 的张量，_n_dims 是多少？这对 compute_offset 的调用有什么影响？

6.3 下节预告

第 4 课 将深入 Memory 类和内存分配器体系：

• Memory 如何封装底层 raw pointer 与分配器
• MemBufferAllocatorBaseObject 接口设计
• CPU 与 CUDA 分配器的具体实现差异
• copy_from 中隐式处理 H2D/D2H 拷贝的设计巧思

---

项目地址：

GitHub：github.com/NKKdev/TFFi... Gitee：gitee.com/NKK_Ovit/tf...

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