Mini-Infer (31): CUDA 后端支持 (下) — TensorRT 风格权重预加载

1. 问题背景：权重拷贝的性能瓶颈

在 GPU 推理中，一个常见的性能问题是权重拷贝开销。

A. PCIe 带宽限制

CPU 和 GPU 之间通过 PCIe 总线通信。即使是 PCIe 4.0 x16，理论带宽也只有约 32 GB/s，远低于 GPU 显存带宽（如 A100 的 2 TB/s）。

1
2
3
CPU Memory ──PCIe──► GPU Memory
              ↑
         瓶颈所在

B. 首次推理延迟

如果每次推理都从 CPU 拷贝权重到 GPU：

1
2
3
推理请求 → 拷贝权重 → 执行计算 → 返回结果
              ↑
         额外延迟

对于一个 100MB 的模型，PCIe 拷贝可能需要 3-5ms，而实际计算可能只需要 1ms。

C. TensorRT 的解决方案

TensorRT 在 Build-Time 将权重加载到 GPU，Run-Time 直接使用：

1
2
Build-Time: 解析模型 → 优化图 → 预加载权重到 GPU → 生成 Engine
Run-Time:   输入数据 → 执行计算 → 返回结果 (权重已在 GPU)

Mini-Infer 借鉴了这一设计。

2. preload_weights_to_gpu() 实现

A. 整体流程

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
// mini_infer/runtime/inference_plan.cpp

core::Status InferencePlan::preload_weights_to_gpu() {
#ifdef MINI_INFER_USE_CUDA
    if (config_.device_type != core::DeviceType::CUDA) {
        return core::Status::SUCCESS;
    }

    MI_LOG_INFO("[InferencePlan] Preloading weights to GPU (TensorRT-style)...");

    size_t total_weight_bytes = 0;
    size_t weight_count = 0;

    // 遍历所有节点
    for (const auto& node : sorted_nodes_) {
        if (!node) continue;

        // 获取常量输入张量（权重/偏置）
        const auto& input_tensors = node->input_tensors();
        for (const auto& tensor : input_tensors) {
            if (!tensor || tensor->empty()) continue;

            // 跳过已在 GPU 的张量
            if (tensor->device() == core::DeviceType::CUDA) continue;

            // 跳过已缓存的张量
            if (gpu_weight_cache_.find(tensor) != gpu_weight_cache_.end()) continue;

            // 创建 GPU 张量并拷贝数据
            auto gpu_tensor = std::make_shared<core::Tensor>(
                tensor->shape(), tensor->dtype(), core::DeviceType::CUDA);

            size_t size_bytes = tensor->size_in_bytes();
            cudaError_t status = cudaMemcpy(
                gpu_tensor->data(), tensor->data(), size_bytes, cudaMemcpyHostToDevice);

            if (status != cudaSuccess) {
                MI_LOG_ERROR("[InferencePlan] Failed to preload weight to GPU");
                return core::Status::ERROR_RUNTIME;
            }

            // 缓存 GPU 张量
            gpu_weight_cache_[tensor] = gpu_tensor;
            total_weight_bytes += size_bytes;
            weight_count++;
        }
    }

    MI_LOG_INFO("[InferencePlan] Preloaded " + std::to_string(weight_count) +
                " weight tensor(s) to GPU (" +
                std::to_string(total_weight_bytes / 1024.0) + " KB total)");

    return core::Status::SUCCESS;
#else
    return core::Status::SUCCESS;
#endif
}

B. 遍历所有节点的 input_tensors

每个节点可能有多个输入：

  • 动态输入：来自上游节点的输出（激活值）。
  • 静态输入：权重、偏置等常量张量。

node->input_tensors() 返回的是静态输入，这些是需要预加载的。

C. 跳过已在 GPU 的 Tensor

1
if (tensor->device() == core::DeviceType::CUDA) continue;

如果模型已经在 GPU 上（例如从 GPU 内存直接加载），就不需要再拷贝。

3. gpu_weight_cache_ 缓存设计

A. 数据结构

1
2
3
4
5
6
// mini_infer/runtime/inference_plan.h

std::unordered_map<std::shared_ptr<const core::Tensor>,
                   std::shared_ptr<core::Tensor>,
                   TensorPtrHash,
                   TensorPtrEqual> gpu_weight_cache_;

Key：CPU 上的原始 Tensor（shared_ptr<const Tensor>）。
Value：GPU 上的对应 Tensor（shared_ptr<Tensor>）。

B. TensorPtrHash 与 TensorPtrEqual

1
2
3
4
5
6
7
8
9
10
11
12
struct TensorPtrHash {
    size_t operator()(const std::shared_ptr<const core::Tensor>& ptr) const {
        return std::hash<const core::Tensor*>{}(ptr.get());
    }
};

struct TensorPtrEqual {
    bool operator()(const std::shared_ptr<const core::Tensor>& lhs,
                    const std::shared_ptr<const core::Tensor>& rhs) const {
        return lhs.get() == rhs.get();
    }
};

为什么用 shared_ptr 作为 key？

  1. 避免悬空指针：如果用裸指针，原始 Tensor 被释放后，key 就变成悬空指针。
  2. 生命周期管理：shared_ptr 确保只要缓存存在，原始 Tensor 就不会被释放。
  3. 唯一性：同一个 Tensor 对象只会被预加载一次。

指针比较 vs 内容比较：

  • 我们比较的是指针，不是内容。
  • 两个内容相同但地址不同的 Tensor 会被视为不同的 key。
  • 这是正确的，因为我们要缓存的是"这个特定的 Tensor 对象"。

4. Run-Time 权重查找

A. get_gpu_tensor() 接口

1
2
3
4
5
6
7
8
9
10
11
12
13
// mini_infer/runtime/inference_plan.cpp

std::shared_ptr<core::Tensor>lan::get_gpu_tensor(
    const std::shared_ptr<const core::Tensor>& cpu_tensor) const {
    if (!cpu_tensor) {
        return nullptr;
    }
    auto it = gpu_weight_cache_.find(cpu_tensor);
    if (it != gpu_weight_cache_.end()) {
        return it->second;
    }
    return nullptr;
}

这个接口供 ExecutionContext 在执行时查找预加载的权重。

B. ensure_on_gpu lambda 的三级查找

在 ExecutionContext::execute_node 中：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
auto ensure_on_gpu = [&](std::shared_ptr<core::Tensor>& tensor) -> core::Status {
    // 0. 已在 GPU 或空张量
    if (!tensor || tensor->device() == core::DeviceType::CUDA) {
        return core::Status::SUCCESS;
    }

    // 第一级：检查 Plan 的预加载缓存 (Build-Time 预加载的权重)
    auto preloaded = plan_->get_gpu_tensor(tensor);
    if (preloaded) {
        tensor = preloaded;
        return core::Status::SUCCESS;
    }

    // 第二级：检查 Context 的本地缓存 (运行时动态创建的张量)
    auto cache_it = gpu_constant_cache_.find(tensor);
    if (cache_it != gpu_constant_cache_.end() && cache_it->second) {
        tensor = cache_it->second;
        return core::Status::SUCCESS;
    }

    // 第三级：运行时拷贝 (应该很少发生)
    auto gpu_tensor = std::make_shared<core::Tensor>(
        tensor->shape(), tensor->dtype(), core::DeviceType::CUDA);
    cudaMemcpy(gpu_tensor->data(), tensor->data(),
               tensor->size_in_bytes(), cudaMemcpyHostToDevice);
    gpu_constant_cache_[tensor] = gpu_tensor;
    tensor = gpu_tensor;
    return core::Status::SUCCESS;
};

三级查找的优先级：

  1. Plan 预加载缓存：Build-Time 预加载的权重，最快。
  2. Context 本地缓存：运行时动态创建的常量，次快。
  3. 运行时拷贝：最后的兜底，最慢。

在正常情况下，所有权重都应该在第一级找到。

5. 输入/输出的 Host-Device 拷贝

权重是静态的，可以预加载。但输入/输出是动态的，每次推理都不同。

A. bind_ordered_inputs 中的 cudaMemcpyAsync

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
// mini_infer/runtime/inference_plan.cpp

core::Status InferencePlan::bind_ordered_inputs(
    ExecutionContext* ctx,
    const std::vector<std::shared_ptr<core::Tensor>>& ordered_inputs) const {

    for (size_t idx = 0; idx < input_bindings_.size(); ++idx) {
        const auto& tensor = ordered_inputs[idx];

#ifdef MINI_INFER_USE_CUDA
        if (config_.device_type == core::DeviceType::CUDA) {
            // 创建或复用 GPU 张量
            if (!outputs[0] || outputs[0]->shape() != tensor->shape() ||
                outputs[0]->device() != core::DeviceType::CUDA) {
                outputs[0] = std::make_shared<core::Tensor>(
                    tensor-tensor->dtype(), core::DeviceType::CUDA);
            }

            // 获取 CUDA Stream
            auto cuda_ctx = std::dynamic_pointer_cast<backends::cuda::CUDADeviceContext>(
                ctx->get_or_create_context(core::DeviceType::CUDA));
            cudaStream_t stream = cuda_ctx ? cuda_ctx->stream() : nullptr;

            // 异步拷贝
            cudaMemcpyAsync(outputs[0]->data(), tensor->data(),
                            tensor->size_in_bytes(),
                            cudaMemcpyHostToDevice, stream);
        }
#endif
    }
    return core::Status::SUCCESS;
}

为什么用 cudaMemcpyAsync？

  • 异步拷贝不会阻塞 CPU。
  • 拷贝和计算可以重叠（如果使用多个 Stream）。
  • 同一 Stream 上的操作会自动同步。

B. collect_ordered_outputs 中的同步拷贝

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
core::Status InferencePlan::collect_ordered_outputs(ExecutionContext* ctx) const {
#ifdef MINI_INFER_USE_CUDA
    if (config_.device_type == core::DeviceType::CUDA && output_tensor) {
        // 同步等待所有 GPU 操作完成
        cudaDeviceSynchronize();

        // 创建 CPU 张量并拷贝
        auto cpu_tensor = std::make_shared<core::Tensor>(
            output_tensor->shape(), output_tensor->dtype());
        cudaMemcpy(cpu_tensor->data(), output_tensor->data(),
                   output_tensor->size_in_bytes(), cudaMemcpyDeviceToHost);
        output_tensor = cpu_tensor;
    }
#endif
    return core::Status::SUCCESS;
}

为什么用同步拷贝？

  • 用户需要立即访问输出数据。
  • cudaDeviceSynchronize 确保所有计算完成。
  • 输出拷贝通常不是瓶颈（输出通常比输入小）。

6. 性能分析与优化空间

A. 当前实现的性能特点

操作 时机 开销
权重预加载 Build-Time 一次性，可接受
输入拷贝 每次推理 异步，可重叠
计算 每次推理 GPU 执行
输出拷贝 每次推理 同步，阻塞

B. 进一步优化方向

  1. 双缓冲 (Double Buffering)：

    • 使用两个输入缓冲区，一个在拷贝，一个在计算。
    • 可以完全隐藏输入拷贝延迟。

  2. Pinned Memory：

    • 使用 cudaMallocHost 分配锁页内存。
    • 可以提高 Host-Device 拷贝速度。

  3. CUDA Graphs：

    • 将整个推理流程录制为 Graph。
    • 减少 Kernel 启动开销。

  4. 多 Stream 并行：

    • 输入拷贝、计算、输出拷贝使用不同 Stream。
    • 实现流水线并行。

7. 总结

本篇我们完成了 CUDA 后端的最后一个关键组件——权重预加载：

  • Build-Time 预加载：preload_weights_to_gpu() 在构建时将权重拷贝到 GPU。
  • gpu_weight_cache_ 缓存：使用 shared_ptr 作为 key，确保生命周期正确。
  • 三级查找机制：Plan 缓存 → Context 缓存 → 运行时拷贝。
  • 异步输入拷贝：使用 cudaMemcpyAsync 减少阻塞。
  • 同步输出拷贝：确保用户能立即访问结果。

至此，Mini-Infer 的 CUDA 后端支持已经完成。下一篇，我们将进入最重要的重构——插件架构，看看如何用 TensorRT 风格的 IPlugin 接口统一算子实现。