C++ 与显存管理：在 AI 推理中实现 PagedAttention 的自定义内存池逻辑

各位同仁、技术爱好者，大家好。

今天，我们将深入探讨在人工智能推理，特别是大型语言模型（LLM）推理中，如何通过C++和CUDA实现精细化的显存管理，以支撑像PagedAttention这样革新性的优化技术。随着AI模型规模的指数级增长，显存已成为推理性能的关键瓶颈。理解并高效利用显存，是构建高性能AI系统的基石。

1. AI推理中的显存挑战：KV Cache与碎片化

在大型语言模型中，自回归（auto-regressive）生成是主流的推理方式。模型在生成每个新token时，需要访问之前所有已生成token的Key（K）和Value（V）表示，这些K和V向量被称为KV Cache。KV Cache的大小与序列长度成正比，通常在生成长序列时占用大量显存。

1.1 KV Cache的存储模式及其低效性

传统的KV Cache管理方式，通常是为每个并发请求预留一块连续的显存区域，足以容纳其最大可能的序列长度。这种方式存在以下问题：

1. 显存碎片化（Fragmentation）：不同请求的序列长度差异巨大。如果一个请求只生成了短序列，其预留的显存大部分是空闲的，但由于是连续分配，这部分空闲显存无法被其他请求利用，导致显存碎片化。
2. 显存利用率低：即使是生成长序列的请求，其KV Cache也并非一开始就完全填充。随着序列的增长，KV Cache逐步填充。在推理初期，大量预留的显存处于空闲状态。
3. 批处理（Batching）受限：由于需要为每个请求预留最大长度的显存，即使总的token数不多，显存也可能很快耗尽，限制了可以同时处理的请求数量（吞吐量）。

考虑一个场景：我们有三个并发请求，最大序列长度为2048。

• 请求A生成了500个token。
• 请求B生成了100个token。
• 请求C生成了1800个token。

如果采用连续分配，每个请求都需要2048长度的KV Cache显存。总显存需求为 $3 times 2048 times (text{key_dim} + text{value_dim}) times text{num_layers} times text{sizeof(float16)}$。其中，请求A和B的显存利用率极低，且其未使用部分无法被其他请求复用。

1.2 PagedAttention：借鉴操作系统内存管理

为了解决这些问题，PagedAttention应运而生。它借鉴了操作系统中的虚拟内存和分页（paging）思想，将KV Cache拆分成固定大小的块（blocks），而不是连续的大块。

核心思想：

• 将每个请求的KV Cache逻辑上划分为固定大小的“逻辑块”（logical blocks）。
• 这些逻辑块可以映射到显存中任意位置的“物理块”（physical blocks）。
• 物理块是固定大小的显存区域。
• 当一个请求需要更多KV Cache时，我们只为其分配一个或多个物理块，而不是扩展其连续区域。

这种机制带来了显著优势：

• 减少碎片化：未使用的KV Cache空间可以以块为单位被其他请求复用，大大减少了内部和外部碎片。
• 提高显存利用率：显存可以更紧凑地被填充，从而在相同的物理显存容量下支持更多的并发请求或更长的序列。
• 支持更大的批处理量：更高的显存利用率直接转化为更高的吞吐量。
• 简化内存管理：固定大小的块使得分配和回收操作更为简单高效。

为了实现PagedAttention，我们需要一套高效的自定义显存管理机制，能够快速分配和回收这些固定大小的物理块。

2. CUDA显存管理基础与性能考量

在深入自定义内存池之前，我们先回顾一下CUDA C++中显存管理的基本操作及其性能特征。

2.1 cudaMalloc 与 cudaFree

这是CUDA中最基本的显存分配和释放函数，对应于CPU上的malloc和free。

#include <cuda_runtime.h>
#include <iostream>

void basicCudaMallocExample() {
    float* device_ptr;
    size_t size = 1024 * sizeof(float); // 1KB

    // 分配显存
    cudaError_t err = cudaMalloc((void**)&device_ptr, size);
    if (err != cudaSuccess) {
        std::cerr << "cudaMalloc failed: " << cudaGetErrorString(err) << std::endl;
        return;
    }
    std::cout << "Allocated " << size << " bytes on device at " << device_ptr << std::endl;

    // ... 使用 device_ptr ...

    // 释放显存
    err = cudaFree(device_ptr);
    if (err != cudaSuccess) {
        std::cerr << "cudaFree failed: " << cudaGetErrorString(err) << std::endl;
        return;
    }
    std::cout << "Freed device memory at " << device_ptr << std::endl;
}

性能考量：

• cudaMalloc和cudaFree是同步（blocking）操作，它们会阻塞CPU线程直到显存操作完成。
• 每次调用都会涉及操作系统级别的内存管理（对于GPU驱动而言），开销相对较高。
• 频繁的小块分配和释放会导致显著的性能开销，并可能加剧显存碎片化。

2.2 统一内存（Unified Memory, UVM）- cudaMallocManaged

统一内存是CUDA 6.0引入的一项功能，它允许CPU和GPU共享同一个虚拟地址空间，并在需要时自动进行数据迁移。

#include <cuda_runtime.h>
#include <iostream>

void unifiedMemoryExample() {
    float* managed_ptr;
    size_t size = 1024 * sizeof(float); // 1KB

    // 分配统一内存
    cudaError_t err = cudaMallocManaged((void**)&managed_ptr, size);
    if (err != cudaSuccess) {
        std::cerr << "cudaMallocManaged failed: " << cudaGetErrorString(err) << std::endl;
        return;
    }
    std::cout << "Allocated " << size << " bytes in unified memory at " << managed_ptr << std::endl;

    // CPU可以直接访问和修改
    managed_ptr[0] = 42.0f;
    std::cout << "CPU set managed_ptr[0] = " << managed_ptr[0] << std::endl;

    // GPU也可以直接访问
    // ... 在CUDA核函数中访问 managed_ptr ...

    // 释放统一内存
    err = cudaFree(managed_ptr);
    if (err != cudaSuccess) {
        std::cerr << "cudaFree failed: " << cudaGetErrorString(err) << std::endl;
        return;
    }
    std::cout << "Freed unified memory at " << managed_ptr << std::endl;
}

性能考量：

• 虽然简化了编程模型，但数据迁移的开销可能很高，尤其是在CPU和GPU频繁交替访问同一区域时。
• 对于高性能计算，通常倾向于显式地使用cudaMalloc分配设备内存，并通过cudaMemcpy进行控制。

2.3 内存对齐

在GPU编程中，内存对齐至关重要。GPU的内存控制器通常以128字节、256字节或512字节的粒度进行数据传输。如果数据没有正确对齐，可能会导致额外的内存事务，降低访存效率。cudaMalloc通常会返回一个256字节对齐的指针，但当我们在自定义内存池中管理子块时，需要特别注意子块的起始地址对齐。

2.4 结论

由于cudaMalloc和cudaFree的同步阻塞特性及较高开销，它们不适用于PagedAttention中频繁的、小块的显存分配和回收。我们需要一个在GPU上运行的、快速的、非阻塞的自定义内存池。

3. 自定义显存池的设计原理

自定义内存池的核心思想是：一次性向系统（CUDA驱动）申请一大块连续的显存，然后在这个大块内部进行细粒度的管理和分配。这样，系统调用开销被平摊，后续的分配和回收操作变得极快。

3.1 内存池的优势

• 极高的分配/回收速度：避免了昂贵的系统调用，操作在用户空间完成。
• 减少碎片化：通过预先规划和固定块大小，可以有效控制碎片。
• 确定性性能：分配和回收的时间复杂度通常是O(1)或O(log N)，而非传统malloc的动态且不可预测。
• 更好的局部性：分配的内存块通常在物理上是相邻的，有助于缓存命中和带宽利用。

3.2 内存池的基本结构

一个典型的固定大小块内存池会包含：

1. 一个大的预分配显存区域：这是所有小块的来源。
2. 一个空闲块列表/位图：用于追踪哪些小块是可用的，哪些正在被使用。
3. 互斥锁（Mutex）：如果内存池可能被多个CPU线程同时访问，需要保证线程安全。

3.3 PagedAttention的内存池需求

对于PagedAttention，我们知道KV Cache块是固定大小的。因此，设计一个固定大小块的内存池是理想选择。

• 块大小：由KV Cache的维度决定，例如 [num_heads, head_dim, block_size]。
• 数量：池中总共可以管理的物理块数量。
• 分配/回收：必须是高效的。

我们将构建一个 GPUMemoryPool 类来管理这些物理块。

4. 实现自定义GPU内存池：GPUMemoryPool

现在，我们开始设计和实现我们的 GPUMemoryPool。

4.1 KV Cache块的结构

首先，定义一个KV Cache块的存储结构。一个块需要存储Key数据和Value数据。

假设：

• max_block_size：一个KV Cache块能够存储的最大token数量。
• num_heads：注意力头的数量。
• head_dim：每个注意力头的维度。
• num_layers：模型层数。
• dtype：数据类型，通常是half（float16）。

一个KV Cache物理块的显存大小为：
block_size_bytes = max_block_size * (num_heads * head_dim * sizeof(dtype_k) + num_heads * head_dim * sizeof(dtype_v)) * num_layers。
通常，K和V的维度和数据类型相同，所以简化为：
block_size_bytes = max_block_size * num_heads * head_dim * sizeof(dtype) * 2 * num_layers。

为了简化，我们通常将num_layers的维度放在外层或在核函数中进行索引。在内存池层面，我们关注的是一个注意力层的一个KV Cache块的存储。
所以一个物理块通常存储max_block_size个token的K和V数据：
block_size_bytes = max_block_size * num_heads * head_dim * sizeof(dtype) * 2。
（假设Key和Value在同一个物理块中，并且K和V紧邻存储。）

例如，如果max_block_size = 16，num_heads = 32，head_dim = 128，dtype = half：
block_size_bytes = 16 * 32 * 128 * sizeof(half) * 2 = 16 * 32 * 128 * 2 * 2 = 262144 bytes = 256 KB。

为了更好的内存访问模式，通常会将K和V数据在块内分开存放，或者将K和V数据分别分配在两个独立的物理块中。这里我们假设K和V数据紧密地存储在一个物理块内。

4.2 GPUMemoryPool 类定义

#pragma once

#include <cuda_runtime.h>
#include <vector>
#include <mutex>
#include <stdexcept>
#include <numeric> // for std::iota

// 辅助函数：检查CUDA错误
#define CUDA_CHECK(call) 
    do { 
        cudaError_t err = call; 
        if (err != cudaSuccess) { 
            fprintf(stderr, "CUDA Error: %s at %s:%dn", cudaGetErrorString(err), __FILE__, __LINE__); 
            throw std::runtime_error(cudaGetErrorString(err)); 
        } 
    } while (0)

class GPUMemoryPool {
public:
    // 构造函数：预分配显存并初始化内存池
    GPUMemoryPool(size_t block_size_bytes, size_t num_blocks, int device_id = 0);

    // 析构函数：释放所有显存
    ~GPUMemoryPool();

    // 分配一个物理块，返回其索引
    // 返回 -1 表示分配失败
    int allocate_block();

    // 释放一个物理块（通过索引）
    void deallocate_block(int block_idx);

    // 获取某个物理块的设备指针
    void* get_block_ptr(int block_idx) const;

    // 获取池中总块数
    size_t get_total_blocks() const { return total_blocks_; }

    // 获取池中空闲块数
    size_t get_free_blocks() const { return free_blocks_count_; }

    // 获取每个块的大小（字节）
    size_t get_block_size_bytes() const { return block_size_bytes_; }

private:
    void* base_ptr_;                 // 内存池的起始设备指针
    size_t block_size_bytes_;        // 每个块的大小（字节）
    size_t total_blocks_;            // 内存池中总块数
    std::vector<int> free_block_indices_; // 存储空闲块索引的列表
    std::mutex mutex_;               // 保护free_block_indices_的互斥锁
    size_t free_blocks_count_;       // 当前空闲块的数量
};

// --- 实现部分 ---

GPUMemoryPool::GPUMemoryPool(size_t block_size_bytes, size_t num_blocks, int device_id)
    : block_size_bytes_(block_size_bytes), total_blocks_(num_blocks), free_blocks_count_(num_blocks) {

    CUDA_CHECK(cudaSetDevice(device_id));

    size_t total_pool_size = block_size_bytes_ * total_blocks_;

    // 1. 预分配一大块GPU显存
    CUDA_CHECK(cudaMalloc(&base_ptr_, total_pool_size));

    // 2. 初始化空闲块列表
    // 初始时，所有块都是空闲的
    free_block_indices_.reserve(total_blocks_);
    for (int i = 0; i < total_blocks_; ++i) {
        free_block_indices_.push_back(i);
    }
    // 或者使用 std::iota (C++11)
    // free_block_indices_.resize(total_blocks_);
    // std::iota(free_block_indices_.begin(), free_block_indices_.end(), 0);

    std::cout << "GPUMemoryPool initialized on device " << device_id 
              << ": " << total_blocks_ << " blocks of " 
              << block_size_bytes_ << " bytes each. Total: " 
              << (double)total_pool_size / (1024 * 1024) << " MB." << std::endl;
}

GPUMemoryPool::~GPUMemoryPool() {
    if (base_ptr_) {
        CUDA_CHECK(cudaFree(base_ptr_));
        base_ptr_ = nullptr;
        std::cout << "GPUMemoryPool freed." << std::endl;
    }
}

int GPUMemoryPool::allocate_block() {
    std::lock_guard<std::mutex> lock(mutex_); // 保护对free_block_indices_的访问

    if (free_block_indices_.empty()) {
        // 没有可用块了
        return -1;
    }

    // 从空闲列表的末尾取出一个块的索引（LIFO策略，简单高效）
    int block_idx = free_block_indices_.back();
    free_block_indices_.pop_back();
    free_blocks_count_--;

    return block_idx;
}

void GPUMemoryPool::deallocate_block(int block_idx) {
    // 确保索引有效
    if (block_idx < 0 || block_idx >= total_blocks_) {
        std::cerr << "Warning: Attempted to deallocate invalid block index: " << block_idx << std::endl;
        return;
    }

    std::lock_guard<std::mutex> lock(mutex_); // 保护对free_block_indices_的访问

    // 检查是否重复释放，防止将同一个块添加多次
    // 这是一个简单的检查，但O(N)复杂度。更好的方法是使用位图或集合。
    // For production, a bitmap or std::vector<bool> is often preferred for O(1) check and better memory.
    // For small total_blocks_ (e.g., thousands), vector search is acceptable.
    // Here, we assume the caller correctly manages block ownership.
    // If block_idx is already in free_block_indices_, it indicates a bug.
    // For simplicity, we just push it back.

    free_block_indices_.push_back(block_idx);
    free_blocks_count_++;
}

void* GPUMemoryPool::get_block_ptr(int block_idx) const {
    if (block_idx < 0 || block_idx >= total_blocks_) {
        throw std::out_of_range("Block index out of range.");
    }
    // 计算物理块的起始地址
    return static_cast<char*>(base_ptr_) + block_idx * block_size_bytes_;
}

代码解释：

• 构造函数：
  • cudaSetDevice(device_id)：指定在哪块GPU上进行操作。
  • cudaMalloc(&base_ptr_, total_pool_size)：一次性分配整个内存池所需的显存。
  • free_block_indices_：使用 std::vector<int> 作为空闲列表。初始时，所有块的索引（0到total_blocks_-1）都被加入到这个列表中。我们使用 std::vector 的 push_back 和 pop_back 来实现LIFO（后进先出）的分配策略，这通常是最快的。
• 析构函数：
  • cudaFree(base_ptr_)：释放预分配的显存。
• allocate_block()：
  • 使用 std::lock_guard<std::mutex> lock(mutex_) 保证线程安全。
  • 从 free_block_indices_ 的末尾取出一个索引，表示分配了一个块。
  • 如果列表为空，则表示显存池已满，返回 -1。
• deallocate_block()：
  • 同样使用 std::lock_guard 保证线程安全。
  • 将 block_idx 放回到 free_block_indices_ 中，表示该块现在可用。
• get_block_ptr()：
  • 通过基地址 base_ptr_ 和块索引 block_idx，计算出对应物理块的实际设备指针。这是在核函数中访问KV Cache的关键。

性能优化和替代方案：

• free_block_indices_ 的替代：
  • std::vector<bool> 或位图（Bitmap）：对于大量的块，位图是更节省内存的选择。每个位代表一个块的空闲状态。分配时查找第一个设置的位，释放时设置对应的位。查找空闲位可能略慢于 std::vector::back()，但内存效率更高，且可以实现O(1)的重复释放检查。
  • 链表：如果需要支持更复杂的分配策略（如最佳适配），链表可能更灵活，但通常引入额外指针开销。
• 线程安全：std::mutex 在CPU侧保护内存池。对于GPU内部的并发访问，这不适用。但通常内存池的分配和释放是由CPU线程发起的。
• 内存对齐：cudaMalloc 保证了 base_ptr_ 是对齐的。由于所有块的大小都相同且是 block_size_bytes_ 的倍数，所以每个 block_idx * block_size_bytes_ 计算出的地址也保持了相同的对齐。确保 block_size_bytes_ 本身是较大的对齐粒度（例如256字节）的倍数，可以进一步优化。

5. PagedAttention上下文与KV Cache管理

有了 GPUMemoryPool，我们还需要一个高级管理器来将逻辑请求映射到物理块。这个管理器将负责：

1. 为每个新的推理请求分配KV Cache块。
2. 在推理过程中，当序列增长需要更多KV Cache时，动态分配新的块。
3. 在请求完成后，释放所有关联的KV Cache块。
4. 提供给CUDA核函数物理块的实际地址列表。

5.1 PagedAttentionManager 类设计

#pragma once

#include "GPUMemoryPool.h"
#include <map>
#include <vector>
#include <list> // For managing request order/eviction policies
#include <numeric> // For std::iota

// 假设的请求ID类型
using RequestId = long long;

// 存储每个逻辑块到物理块的映射
// 实际的PagedAttention会有一个更复杂的 block_table
// 这里简化为每个请求持有的物理块索引列表
struct RequestKVInfo {
    std::vector<int> physical_block_indices; // 存储该请求已分配的物理块索引
    // 其他请求相关信息，如：
    // int sequence_length; // 当前序列长度
    // int current_token_count; // 已存储的token数量
    // int max_total_tokens; // 最大允许的总token数量
};

class PagedAttentionManager {
public:
    // 构造函数
    PagedAttentionManager(size_t block_size_bytes, size_t num_pool_blocks, int device_id);

    // 析构函数
    ~PagedAttentionManager();

    // 为一个新请求分配初始的KV Cache块
    // 返回true表示成功，false表示显存不足
    bool allocate_kv_cache_for_request(RequestId request_id, int num_initial_blocks);

    // 为一个请求动态增加一个KV Cache块
    // 返回新分配块的物理索引，-1表示失败
    int add_block_to_request(RequestId request_id);

    // 释放一个请求的所有KV Cache块
    void free_kv_cache_for_request(RequestId request_id);

    // 获取某个请求的所有物理块的设备指针
    // 返回一个设备指针数组，供核函数使用
    // 注意：这个数组本身需要分配在设备内存上
    void** get_physical_block_pointers_device(RequestId request_id);

    // 获取某个请求的当前物理块索引列表（用于调试或高级调度）
    const std::vector<int>& get_request_physical_block_indices(RequestId request_id) const;

    // 获取内存池信息
    size_t get_free_blocks_in_pool() const { return pool_.get_free_blocks(); }
    size_t get_total_blocks_in_pool() const { return pool_.get_total_blocks(); }

private:
    GPUMemoryPool pool_; // 底层GPU内存池
    std::map<RequestId, RequestKVInfo> request_kv_cache_map_; // 映射请求ID到其KV信息
    std::mutex manager_mutex_; // 保护request_kv_cache_map_的互斥锁

    // 缓存设备端指针数组，避免每次都重新分配和拷贝
    void** device_block_pointers_cache_; // 设备端指针数组的基地址
    std::vector<void*> host_block_pointers_staging_; // 宿主端用于staging的指针数组

    // 最大物理块数量（与pool_的total_blocks_相同）
    size_t max_physical_blocks_;

    // 辅助函数：将host_block_pointers_staging_拷贝到device_block_pointers_cache_
    void update_device_block_pointers(RequestId request_id, const std::vector<int>& physical_indices);
};

// --- 实现部分 ---

PagedAttentionManager::PagedAttentionManager(size_t block_size_bytes, size_t num_pool_blocks, int device_id)
    : pool_(block_size_bytes, num_pool_blocks, device_id), max_physical_blocks_(num_pool_blocks) {

    // 预分配设备端指针数组，用于传递给核函数
    // 最大可能需要传递的指针数量是池中所有物理块的数量
    CUDA_CHECK(cudaMalloc(&device_block_pointers_cache_, sizeof(void*) * max_physical_blocks_));
    host_block_pointers_staging_.resize(max_physical_blocks_);
    std::cout << "PagedAttentionManager initialized. Device pointer cache allocated for " 
              << max_physical_blocks_ << " blocks." << std::endl;
}

PagedAttentionManager::~PagedAttentionManager() {
    // 确保所有请求的KV Cache都被释放
    for (auto const& [request_id, kv_info] : request_kv_cache_map_) {
        // 实际应用中可能需要更优雅的错误处理，或者确保在此之前所有请求都已完成并释放
        std::cerr << "Warning: Request " << request_id << " still has KV cache allocated during manager destruction." << std::endl;
        // 这里只是打印警告，pool_的析构函数会处理底层cudaFree
    }
    if (device_block_pointers_cache_) {
        CUDA_CHECK(cudaFree(device_block_pointers_cache_));
        device_block_pointers_cache_ = nullptr;
    }
    std::cout << "PagedAttentionManager freed." << std::endl;
}

bool PagedAttentionManager::allocate_kv_cache_for_request(RequestId request_id, int num_initial_blocks) {
    std::lock_guard<std::mutex> lock(manager_mutex_);

    if (request_kv_cache_map_.count(request_id)) {
        std::cerr << "Error: Request ID " << request_id << " already exists." << std::endl;
        return false;
    }

    RequestKVInfo kv_info;
    kv_info.physical_block_indices.reserve(num_initial_blocks); // 预留空间

    for (int i = 0; i < num_initial_blocks; ++i) {
        int block_idx = pool_.allocate_block();
        if (block_idx == -1) {
            // 分配失败，回滚已分配的块
            for (int allocated_idx : kv_info.physical_block_indices) {
                pool_.deallocate_block(allocated_idx);
            }
            std::cerr << "Failed to allocate " << num_initial_blocks << " initial blocks for request " 
                      << request_id << ". Not enough memory." << std::endl;
            return false;
        }
        kv_info.physical_block_indices.push_back(block_idx);
    }

    request_kv_cache_map_[request_id] = kv_info;
    std::cout << "Allocated " << num_initial_blocks << " blocks for request " << request_id << std::endl;
    return true;
}

int PagedAttentionManager::add_block_to_request(RequestId request_id) {
    std::lock_guard<std::mutex> lock(manager_mutex_);

    auto it = request_kv_cache_map_.find(request_id);
    if (it == request_kv_cache_map_.end()) {
        std::cerr << "Error: Request ID " << request_id << " not found." << std::endl;
        return -1;
    }

    int block_idx = pool_.allocate_block();
    if (block_idx == -1) {
        std::cerr << "Failed to add block for request " << request_id << ". Not enough memory." << std::endl;
        return -1;
    }

    it->second.physical_block_indices.push_back(block_idx);
    std::cout << "Added block " << block_idx << " to request " << request_id << ". Total blocks: " 
              << it->second.physical_block_indices.size() << std::endl;
    return block_idx;
}

void PagedAttentionManager::free_kv_cache_for_request(RequestId request_id) {
    std::lock_guard<std::mutex> lock(manager_mutex_);

    auto it = request_kv_cache_map_.find(request_id);
    if (it == request_kv_cache_map_.end()) {
        std::cerr << "Warning: Attempted to free non-existent request ID: " << request_id << std::endl;
        return;
    }

    for (int block_idx : it->second.physical_block_indices) {
        pool_.deallocate_block(block_idx);
    }
    request_kv_cache_map_.erase(it);
    std::cout << "Freed all blocks for request " << request_id << std::endl;
}

const std::vector<int>& PagedAttentionManager::get_request_physical_block_indices(RequestId request_id) const {
    std::lock_guard<std::mutex> lock(manager_mutex_);
    auto it = request_kv_cache_map_.find(request_id);
    if (it == request_kv_cache_map_.end()) {
        throw std::runtime_error("Request ID not found for getting physical block indices.");
    }
    return it->second.physical_block_indices;
}

// 辅助函数实现
void PagedAttentionManager::update_device_block_pointers(RequestId request_id, const std::vector<int>& physical_indices) {
    if (physical_indices.empty()) {
        // 如果没有块，则无需更新
        return;
    }

    // 1. 在host_block_pointers_staging_中填充实际的设备指针
    for (size_t i = 0; i < physical_indices.size(); ++i) {
        host_block_pointers_staging_[i] = pool_.get_block_ptr(physical_indices[i]);
    }

    // 2. 将这些指针从Host拷贝到Device
    CUDA_CHECK(cudaMemcpy(device_block_pointers_cache_, 
                          host_block_pointers_staging_.data(), 
                          physical_indices.size() * sizeof(void*), 
                          cudaMemcpyHostToDevice));
}

void** PagedAttentionManager::get_physical_block_pointers_device(RequestId request_id) {
    std::lock_guard<std::mutex> lock(manager_mutex_);
    auto it = request_kv_cache_map_.find(request_id);
    if (it == request_kv_cache_map_.end()) {
        throw std::runtime_error("Request ID not found for getting device block pointers.");
    }

    // 每次需要时，根据当前请求的物理块索引，更新设备端的指针数组
    update_device_block_pointers(request_id, it->second.physical_block_indices);

    // 返回设备端指针数组的基地址
    return device_block_pointers_cache_;
}

代码解释：

• RequestKVInfo：这是一个简单的结构，用于存储每个请求已分配的物理块索引列表。在实际的PagedAttention实现中，这个结构会更复杂，包含逻辑块到物理块的映射表（例如 std::map<int, int>），以及当前序列长度、注意力掩码等信息。
• request_kv_cache_map_：std::map<RequestId, RequestKVInfo> 将唯一的请求ID映射到其对应的KV Cache信息。
• allocate_kv_cache_for_request()：
  • 为新请求分配指定数量的初始块。
  • 如果显存不足，会回滚已分配的块，确保原子性。
• add_block_to_request()：
  • 在推理过程中，当序列增长需要新块时调用。
• free_kv_cache_for_request()：
  • 请求完成后，遍历其所有物理块索引，并将其返还给 GPUMemoryPool。
• get_physical_block_pointers_device()：
  • 关键功能：这个函数是连接CPU管理逻辑和GPU核函数执行的关键。
  • 它首先从 GPUMemoryPool 获取每个物理块的实际设备指针。
  • 然后，它将这些 void* 指针从Host内存 (host_block_pointers_staging_) 拷贝到预先分配的Device内存 (device_block_pointers_cache_)。
  • 最后，返回 device_block_pointers_cache_ 的基地址。CUDA核函数将接收这个设备指针数组，并根据需要索引其中的元素来访问具体的KV Cache块。

关于 get_physical_block_pointers_device 的重要说明：

每次调用 get_physical_block_pointers_device 都会进行一次 cudaMemcpyHostToDevice 操作。如果这个函数在每个token生成步骤中都被调用，并且 physical_block_indices 长度较长，那么 cudaMemcpy 的开销可能会成为瓶颈。

优化方案：

• 异步拷贝：使用CUDA Stream进行异步拷贝，与核函数执行并行。
• 增量更新：只拷贝新添加的块的指针，而不是整个列表。
• Unified Memory (UVM)：如果 host_block_pointers_staging_ 和 device_block_pointers_cache_ 都使用UVM，可以减少显式拷贝的需要，但仍需注意页迁移开销。
• 设备端调度器：更高级的PagedAttention实现（如vLLM）可能将调度逻辑和块映射表完全放在设备端，从而避免频繁的Host-Device通信。这意味着 request_kv_cache_map_ 甚至可以存储在设备端。

6. CUDA核函数集成

在CUDA核函数中，我们如何使用这些物理块指针来访问KV Cache数据？

假设我们的注意力核函数需要访问某个请求的KV Cache。它将接收一个设备指针数组 kv_cache_block_ptrs 和一个映射表 block_table。

block_table 是一个设备端数组，表示一个请求的逻辑块索引到物理块索引的映射。
例如，block_table[logical_block_idx] 给出 physical_block_idx。

#include <cuda_fp16.h> // For half precision

// 假设的注意力核函数
__global__ void paged_attention_kernel(
    half* output,              // 输出tensor
    const half* query,         // Q tensor
    const void** kv_cache_block_ptrs, // 设备端指针数组，每个元素指向一个物理KV Cache块
    const int* block_table,     // 设备端块映射表：logical_block_idx -> physical_block_idx
    int num_seqs_in_batch,     // 当前批次中的序列数量
    int max_seq_len,           // 批次中最大序列长度
    int block_size,            // 每个KV Cache块存储的token数量
    int num_heads,
    int head_dim,
    int num_layers,             // 模型的层数
    int current_layer_idx       // 当前计算的注意力层索引
) {
    // 假设每个线程处理一个输出token或一个注意力头
    // 简化起见，我们只展示如何获取KV数据
    int sequence_idx = blockIdx.x; // 批次中的序列索引
    int token_idx_in_seq = threadIdx.x; // 序列中的token索引 (假设每个线程处理一个token)

    if (sequence_idx >= num_seqs_in_batch) return;
    if (token_idx_in_seq >= max_seq_len) return; // 实际应是当前序列的长度

    // 计算当前token对应的逻辑块索引和块内偏移
    int logical_block_idx = token_idx_in_seq / block_size;
    int token_offset_in_block = token_idx_in_seq % block_size;

    // 获取当前序列对应的物理块索引
    // 注意：这里的block_table需要根据请求ID和当前层的逻辑来正确索引
    // 实际的block_table可能是一个二维数组或更复杂的结构，
    // 例如 block_table[request_idx * max_num_logical_blocks_per_seq + logical_block_idx]
    // 简化处理，假设block_table是针对当前序列的
    int physical_block_idx = block_table[logical_block_idx];

    // 从设备指针数组中获取该物理块的基地址
    // 注意：kv_cache_block_ptrs[physical_block_idx] 是一个void*，需要转换为实际类型
    char* kv_block_base_ptr = (char*)kv_cache_block_ptrs[physical_block_idx];

    // 在物理块内部，Key和Value数据需要进一步索引
    // 假设一个物理块存储了 `block_size` 个token的K和V数据，
    // 并且K数据在前，V数据在后，每个token的K和V数据都是 `num_heads * head_dim` 大小。
    size_t head_data_size = num_heads * head_dim * sizeof(half); // 一个token的K或V数据大小

    // Key数据偏移：
    // 每个物理块内存储了所有层的KV数据。
    // 如果KV Cache是按层划分的，则 kv_block_base_ptr 已经是针对当前层的。
    // 如果一个物理块存储了所有层的KV数据，那么还需要考虑层偏移。
    // 假设 kv_block_base_ptr 已经指向当前层的KV数据区域
    // K数据通常位于块的前半部分，V数据位于后半部分
    size_t kv_half_block_size_bytes = block_size * head_data_size; // K或V各自的区域大小

    half* k_data_ptr = (half*)(kv_block_base_ptr + current_layer_idx * (2 * kv_half_block_size_bytes) + 0 * kv_half_block_size_bytes);
    half* v_data_ptr = (half*)(kv_block_base_ptr + current_layer_idx * (2 * kv_half_block_size_bytes) + 1 * kv_half_block_size_bytes);

    // 访问特定token的Key和Value
    half* current_k_ptr = k_data_ptr + token_offset_in_block * num_heads * head_dim;
    half* current_v_ptr = v_data_ptr + token_offset_in_block * num_heads * head_dim;

    // ... 使用 current_k_ptr 和 current_v_ptr 进行注意力计算 ...
    // 例如：
    // half key_value = current_k_ptr[some_head_idx * head_dim + some_dim_idx];
}

核函数集成要点：

• kv_cache_block_ptrs：这是由 PagedAttentionManager::get_physical_block_pointers_device() 返回的设备端指针数组。核函数通过 kv_cache_block_ptrs[physical_block_idx] 获取指定物理块的基地址。
• block_table：这是一个额外的设备端数组，由CPU侧的 PagedAttentionManager 或更上层的调度器构建，用于将一个请求的逻辑序列位置映射到其对应的物理块索引。例如，对于序列中的第 i 个token，其逻辑块索引是 i / block_size。然后通过 block_table[i / block_size] 找到对应的 physical_block_idx。
• 内存布局：KV Cache块内部的内存布局需要严格定义。例如，K和V数据是否在一个块中？K和V数据是否按层存储？这些都会影响核函数中指针的计算。这里假设K和V在一个块中，且K在前V在后，并考虑了层索引。

7. 高级考量与进一步优化

7.1 多GPU支持

如果AI模型部署在多GPU环境中，每个GPU都需要维护自己的 GPUMemoryPool。PagedAttentionManager 可以扩展为管理多个 GPUMemoryPool 实例，并根据请求的分配策略（例如，将KV Cache均匀分布到所有GPU，或者根据模型层数将其分配到不同的GPU上）来决定向哪个 GPUMemoryPool 申请显存。

7.2 异步操作与CUDA Stream

cudaMemcpyHostToDevice 通常是同步操作，但我们可以利用CUDA Stream实现异步拷贝。

// 在 PagedAttentionManager 中
// CUDA_CHECK(cudaMemcpyAsync(device_block_pointers_cache_, 
//                           host_block_pointers_staging_.data(), 
//                           physical_indices.size() * sizeof(void*), 
//                           cudaMemcpyHostToDevice, 
//                           stream)); // 传入一个CUDA Stream

核函数的启动也可以绑定到同一个Stream，以实现拷贝和计算的重叠。

7.3 动态池大小调整

当前的内存池是固定大小的。在某些场景下，可能需要动态调整池的大小（例如，在显存压力较小的时候缩小，或在显存不足时尝试扩大）。这会增加复杂性：

• 扩大：需要重新 cudaMalloc 一个更大的区域，并将旧数据拷贝过来，然后 cudaFree 旧区域。这通常是昂贵的。
• 缩小：释放一部分末尾的块，但需要确保这些块都是空闲的。

通常，为了性能和稳定性，内存池在启动时就被配置为最大可能大小。

7.4 调度与驱逐策略

当显存池满时，allocate_block() 会失败。在生产环境中，需要实现一套调度和驱逐（eviction）策略：

• LRU（最近最少使用）：驱逐最长时间未被访问的请求。
• LFU（最不常用）：驱逐使用频率最低的请求。
• 优先级调度：根据请求的优先级或SLA（服务等级协议）来决定驱逐哪些请求。
• 预填充/预分配：在空闲时段预先分配一些块给可能到来的请求。

这些策略通常在 PagedAttentionManager 的上层调度器中实现。

7.5 内存监控与调试

在复杂的显存管理系统中，监控内存使用情况至关重要：

• cudaMemGetInfo：获取GPU总显存和可用显存。
• 自定义统计：在 GPUMemoryPool 中增加已用块/空闲块的计数器。
• CUDA Nsight Systems：强大的性能分析工具，可以可视化CUDA API调用、核函数执行和内存操作，帮助识别瓶颈。

8. 性能考量与总结展望

通过实现自定义的显存池和PagedAttention机制，我们能够显著提升LLM推理的性能和显存利用率：

• 吞吐量提升：由于碎片化减少，可以在相同显存下支持更多的并发请求，从而提高批处理量和QPS（Queries Per Second）。
• 延迟降低：高效的块分配和回收减少了内存管理的开销。
• 显存利用率优化：减少了KV Cache的闲置空间，使得宝贵的GPU显存能够被更有效地利用。

特性	传统连续KV Cache分配	PagedAttention + 自定义内存池
内存碎片化	严重	极低
显存利用率	低	高
分配/释放速度	慢 (cudaMalloc/Free)	快 (内存池内部操作)
批处理量	受限于最大序列长度	受限于总物理块数
调度复杂度	相对简单	需管理逻辑-物理块映射、驱逐策略

自定义内存池的实现虽然增加了初始的开发复杂性，但它为高性能AI推理提供了坚实的基础。这是C++和CUDA深度结合，精细化显存管理的典型案例，也是构建高效AI基础设施不可或缺的一环。随着AI模型向万亿参数规模迈进，对底层计算资源的精细控制将变得越来越重要。理解并掌握这些技术，是每一位AI系统工程师的必备技能。