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

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

6. 关键技术实现

6.1 Q8_0 量化

量化格式:

每 32 个元素为一组:
- 32 × int8 (权重)
- 1 × fp16 (scale)
- 1 × fp16 (可选 zero point)

实际存储: 32 bytes (weights) + 2 bytes (scale) = 34 bytes
压缩率: 34 / (32 * 2) = 53.125% (相比 FP16)

量化 Kernel:

__global__ void quant_q8_0_kernel(const float* src, int8_t* dst, 
                                   half* scales, int num_elements) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int group_idx = idx / 32;
    int intra_group = idx % 32;
  
    // 1. 计算每组最大值
    __shared__ float max_val[256];
    float abs_val = fabs(src[idx]);
    max_val[threadIdx.x] = abs_val;
    __syncthreads();
  
    // Reduction 找最大值
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (threadIdx.x < stride) {
            max_val[threadIdx.x] = max(max_val[threadIdx.x], 
                                       max_val[threadIdx.x + stride]);
        }
        __syncthreads();
    }
  
    // 2. 计算 scale
    if (intra_group == 0) {
        float max = max_val[0];
        scales[group_idx] = __float2half(max / 127.0f);
    }
    __syncthreads();
  
    // 3. 量化
    float scale = __half2float(scales[group_idx]);
    float scaled = src[idx] / scale;
    dst[idx] = __float2int_rn(scaled);
}

量化矩阵乘法:

__global__ void quant_mat_mul_kernel(const int8_t* A, const int8_t* B,
                                     const half* scales_a, const half* scales_b,
                                     float* C, int M, int N, int K) {
    // 使用 WMMA 或 IMMA 指令
    // 累加时使用 FP32 精度
    // 最后乘以 scale_a * scale_b
}

6.2 RoPE (Rotary Positional Embedding)

数学公式:

RoPE(x, pos) = x · R(pos)

其中 R(pos) 是旋转矩阵:
R(pos) = diag(R_0(pos), R_1(pos), ..., R_{d/2}(pos))

R_i(pos) = [cos(i·θ·pos)  -sin(i·θ·pos)]
           [sin(i·θ·pos)   cos(i·θ·pos)]

θ = 10000^(-2i/d)

预计算优化:

// 预计算 cos/sin 表
void precompute_rope_tables(int max_seq_len, int dim, float base) {
    for (int pos = 0; pos < max_seq_len; ++pos) {
        for (int i = 0; i < dim / 2; ++i) {
            float freq = 1.0f / pow(base, 2.0f * i / dim);
            rope_table[pos * dim + 2 * i] = cos(pos * freq);
            rope_table[pos * dim + 2 * i + 1] = sin(pos * freq);
        }
    }
}

Kernel 实现:

__device__ half2 complex_mul_half2(half2 a, half2 b) {
    // (a.x + a.y*i) * (b.x + b.y*i)
    // = (a.x*b.x - a.y*b.y) + (a.x*b.y + a.y*b.x)*i
    float ax = __half2float(a.x);
    float ay = __half2float(a.y);
    float bx = __half2float(b.x);
    float by = __half2float(b.y);
  
    return __float22half2_rn(make_float2(
        ax * bx - ay * by,
        ax * by + ay * bx
    ));
}

__global__ void rope_kernel(half* q, half* k, 
                           const float* cos_sin_table,
                           int pos, int dim) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int head = idx / dim;
    int d = idx % dim;
  
    if (d < dim / 2) {
        half2 q_val = reinterpret_cast<half2*>(q)[idx];
        half2 cos_sin = __float22half2_rn(
            make_float2(cos_sin_table[pos * dim + 2 * d],
                       cos_sin_table[pos * dim + 2 * d + 1]));
    
        reinterpret_cast<half2*>(q)[idx] = complex_mul_half2(q_val, cos_sin);
    }
}

6.3 Flash Attention v2

核心思想:

Tiling: 分块计算,适应 SRAM
Recomputation: 重计算替代存储
Online Softmax: 在线归一化

算法伪代码:

FlashAttention(Q, K, V):
    split Q, K, V into tiles
  
    for each tile Q_i:
        m_i = -inf, l_i = 0, O_i = 0
    
        for each tile K_j, V_j:
            S_ij = Q_i @ K_j.T / sqrt(d)
        
            m_ij = max(m_i, rowmax(S_ij))
            P_ij = exp(S_ij - m_ij)
            l_ij = l_i * exp(m_i - m_ij) + rowsum(P_ij)
        
            O_i = O_i * (l_i * exp(m_i - m_ij) / l_ij) + 
                  P_ij @ V_j * (exp(m_i - m_ij) / l_ij)
        
            m_i = m_ij
            l_i = l_ij
    
        O_i = O_i / l_i
  
    return O

CUDA 实现要点:

Shared Memory Layout: K/V 分块加载
Register Accumulation: O/m/l 存储在寄存器
Async Copy: Hopper 架构异步拷贝
Swizzle: 避免 Shared Memory Bank Conflict
Fused RoPE: 加载时直接应用 RoPE

6.4 Graph Optimization

6.4.1 算子融合 (Operator Fusion)

融合规则:

Element-wise Ops 融合:
Add + Mul → FusedAddMul
Norm + Mul → FusedNormMul
Reshape + Transpose → FusedReshapeTranspose

Vertical Fusion:
MatMul + Add → FusedGEMM+Bias
MatMul + Norm → FusedGEMM+Norm

Horizontal Fusion:
Multiple Element-wise on same input → Single Kernel

实现:

void GraphOptimizer::fuse_operators(Graph& graph) {
    for (auto& node : graph.nodes()) {
        // 检查是否可以与前驱节点融合
        if (can_fuse_with_predecessor(node)) {
            auto fused_node = create_fused_node(node, node->predecessor());
            graph.replace_nodes({node, node->predecessor()}, fused_node);
        }
    }
}

6.4.2 死代码消除 (Dead Code Elimination)

void GraphOptimizer::eliminate_dead_code(Graph& graph) {
    // 1. 标记所有可达节点
    std::unordered_set<std::shared_ptr<GraphNode>> reachable;
    mark_reachable(graph.output(), reachable);
  
    // 2. 删除不可达节点
    for (auto it = graph.nodes().begin(); it != graph.nodes().end();) {
        if (!reachable.count(*it)) {
            it = graph.nodes().erase(it);
        } else {
            ++it;
        }
    }
}

7. 性能优化策略

7.1 CUDA Kernel 优化

7.1.1 内存访问优化

Coalesced Access: 确保线程束内连续访问
Shared Memory Tiling: 减少全局内存访问
Swizzle: 避免 Bank Conflict
Async Copy: Hopper 架构异步拷贝

7.1.2 计算优化

WMMA/IMMA: Tensor Core 加速矩阵乘法
Loop Unrolling: #pragma unroll
Instruction Level Parallelism: 隐藏延迟
Fused Operations: RoPE + Load, Softmax + Scale

7.1.3 Occupancy 优化

Register Pressure: 控制寄存器使用
Shared Memory: 合理分配共享内存
Thread Block Size: 选择合适的块大小

7.2 显存优化

7.2.1 内存池

预分配大块显存
避免频繁 malloc/free
256 字节对齐

7.2.2 生命周期分析

精确计算张量活跃区间
非重叠区间复用内存
降低峰值显存占用 30-50%

7.2.3 延迟回收

CUDA Event 异步通知
避免阻塞计算流
提高并发度

7.3 计算图优化

7.3.1 算子融合

减少 Kernel Launch 开销
提高数据局部性
降低显存带宽压力

7.3.2 常量折叠

编译期计算常量表达式
减少运行时计算

7.3.3 公共子表达式消除

识别重复计算
缓存中间结果

7.4 KV Cache 优化

7.4.1 Paged Attention

消除内存碎片
支持可变长度
提高内存利用率 20-30%

7.4.2 Prefix Cache (未来)

相同 prompt 复用 KV
减少重复计算

7.4.3 Quantized KV (未来)

INT8/FP8 KV Cache
降低显存占用 50%

7.5 批处理优化

7.5.1 Static Batching

固定 batch size
最大化 GPU 利用率

7.5.2 Continuous Batching (未来)

动态 batch 大小
请求级调度
提高吞吐量 2-5x

8. 扩展性与可维护性

8.1 如何添加新模型

步骤:

创建模型检测器:

class MyModelDetector : public ModelDetectorBase {
    bool matches(const std::string& format) const override {
        return format == "mymodel";
    }
  
    std::shared_ptr<ModelLoaderBase> create_loader() override {
        return std::make_shared<MyModelLoader>();
    }
};

REGISTER_MODULE_OBJECT(MyModelDetector, ModelDetectorBase, 
                      MODEL_DETECTOR_FLAG, "mymodel");

创建模型加载器:

class MyModelLoader : public ModelLoaderBase {
    ModelLoadResult load_from_file(...) override {
        // 解析模型格式
        // 读取元数据
        // 构建权重映射
    }
};

创建模型创建器:

class MyModelCreator : public ModelCreatorBase {
    void build_graph(...) override {
        // 定义模型结构
        // 构建计算图
    }
};

注册到工厂:

REGISTER_MODULE_OBJECT(MyModelCreator, ModelCreatorBase,
                      MODEL_CREATOR_FLAG, "mymodel_arch");

8.2 如何添加新算子

步骤:

定义算子类:

template<typename T, typename DeviceTag>
class MyCustomOp;

创建参数构建器:

class MyCustomOpBuilder : public OpParamBuilderBase<MyCustomOpBuilder> {
    struct Params {
        static constexpr const char* Input = "input";
        static constexpr const char* Output = "output";
    };
  
public:
    template<typename T>
    MyCustomOpBuilder& input(T&& value) {
        return set(Params::Input, value);
    }
};

实现 CUDA Kernel:

__global__ void my_custom_op_kernel(...) {
    // Kernel 实现
}

注册到 FunctionFactory:

FunctionFactory::instance()->register_callback(
    "kernel", "my_custom_op_cuda",
    [](...) { my_custom_op_kernel<<<...>>>(...); }
);

8.3 如何添加新设备后端

步骤:

创建设备类:

class DeviceNewBackend : public DeviceBaseObject {
    void device_init() override {
        // 初始化设备
    }
  
    std::shared_ptr<MemBufferAllocatorBaseObject> 
    get_device_buffer_allocator(int device_id) override {
        return std::make_shared<MemBufferAllocatorNew>();
    }
};

创建内存分配器:

class MemBufferAllocatorNew : public MemBufferAllocatorBaseObject {
    void* allocate(size_t size) override {
        // 新后端的内存分配
    }
  
    void release(void* ptr) override {
        // 新后端的内存释放
    }
};

注册到工厂:

REGISTER_MODULE_OBJECT(DeviceNewBackend, DeviceBaseObject,
                      DEVICE_BACKEND_FLAG, DEVICE_BACKEND_TYPE_NEW);

注：目前更新的项目解析是概览版的，给感兴趣的小伙伴一个大概的项目轮廓。后面会根据内容拆分细节分享。欢迎关注。

附录

A. 关键文件索引

模块	关键文件	行数
运行时	`src/core/runtime/InferRuntime.h`	269
计算图	`src/core/graph/Graph.h`	168
张量	`src/core/mem/Tensor.h`	485
KV Cache	`src/core/runtime/KVCache.h`	476
内存管理	`src/core/runtime/MemManager.h`	277
任务调度	`src/taskgraph/include/TaskFlowSchedule.h`	124
算子创建	`src/kernel/include/TFFOPCreator.h`	1645
Flash Attention	`src/kernel/cuda/flash_attention.cu`	822
模块工厂	`include/ModuleFactory.h`	331

B. 联系方式

项目地址:
- Gitee: gitee.com/NKK_Ovit/tf…
- GitHub: github.com/NKKdev/TFFi…

文档版本: v1.0 最后更新: 2026-04-19 适用版本: TFFInfer v0.x

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

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

目录

6. 关键技术实现

6.1 Q8_0 量化

6.2 RoPE (Rotary Positional Embedding)

6.3 Flash Attention v2

6.4 Graph Optimization

6.4.1 算子融合 (Operator Fusion)

6.4.2 死代码消除 (Dead Code Elimination)

7. 性能优化策略

7.1 CUDA Kernel 优化

7.1.1 内存访问优化

7.1.2 计算优化

7.1.3 Occupancy 优化

7.2 显存优化

7.2.1 内存池

7.2.2 生命周期分析

7.2.3 延迟回收

7.3 计算图优化

7.3.1 算子融合

7.3.2 常量折叠

7.3.3 公共子表达式消除

7.4 KV Cache 优化

7.4.1 Paged Attention

7.4.2 Prefix Cache (未来)

7.4.3 Quantized KV (未来)

7.5 批处理优化

7.5.1 Static Batching

7.5.2 Continuous Batching (未来)

8. 扩展性与可维护性

8.1 如何添加新模型

8.2 如何添加新算子

8.3 如何添加新设备后端

注：目前更新的项目解析是概览版的，给感兴趣的小伙伴一个大概的项目轮廓。后面会根据内容拆分细节分享。欢迎关注。

附录

A. 关键文件索引

B. 联系方式