CUDA Graph 动态形状捕获与 LLM 推理内存池优化深度实战

现代 LLM 推理引擎正面临一个核心矛盾：GPU 计算密度持续攀升，而 CPU 端的 kernel launch 开销成为瓶颈。CUDA Graph 通过捕获并重放 GPU 指令序列，将多次 kernel 调用压缩为一次提交，显著降低了 launch overhead。但 LLM 推理中序列长度动态变化的特点，让 CUDA Graph 的"静态形状"假设面临严峻挑战。本文将深入剖析 CUDA Graph 的捕获机制、动态形状问题的本质生产级优化方案，并给出可落地的内存池设计。

一、为什么 LLM 推理需要 CUDA Graph

在典型的自回归生成（autoregressive decoding）中，每个 token 的生成都需要执行一次完整的模型前向传播。以 LLaMA-2-70B 为例，单次前向涉及 ~80 个 kernel（注意力、FFN、LayerNorm、Resid ual Add 等），当 batch size 为 32、生成 100 个 token 时，总共需要 launch 约 256,000 个 kernel。

传统执行模式下，每个 kernel 的 launch 需要 CPU 将指令写入 GPU 的 work queue，单次 launch 的 CPU 开销约 5-10μs。在 decoding 阶段，由于每次只处理一个 token，GPU 执行时间可能只有几十微秒，CPU launch overhead 占比可达 15%-30%。

CUDA Graph 的核心思路：将一系列 CUDA 调用（kernel launches、memcpy、events）录制为一个有向无环图（DAG），然后通过 cudaGraphLaunch 一次性提交整个图到 GPU。这消除了逐 kernel 的 CPU-GPU 交互，将 launch overhead 从 O(N) 降至 O(1)。

二、CUDA Graph 捕获机制深度解析

2.1 基本捕获流程

CUDA Graph 的捕获分三个阶段：

// 1. 开始捕获
cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal);

// 2. 执行需要录制的操作
for (int i = 0; i < num_layers; i++) {
    launch_attention_kernel<<<grid, block, 0, stream>>>(...);
    launch_ffn_kernel<<<grid, block, 0, stream>>>(...);
    launch_layernorm_kernel<<<grid, block, 0, stream>>>(...);
}

// 3. 结束捕获并实例化
cudaStreamEndCapture(stream, &graph);
cudaGraphInstantiate(&graphExec, graph, NULL, NULL, 0);

关键点在于：捕获期间，所有 memory allocation 操作会被"冻结"——CUDA 驱动记录下分配的大小、地址和生命周期，回放时利用相同的地址布局。

2.2 Graph 内存模型

CUDA Graph 录制的不仅是 kernel，还包括：

• Node dependencies：kernel 之间的依赖关系（由 stream ordering 决定）

• Memory nodes：cudaMalloc、cudaFree、cudaMemcpy 等内存操作

• Host nodes：需要 CPU 回调的操作

• Event nodes：跨 stream 同步事件

实例化（instantiation）时，驱动会：

1. 分配所有需要的 device memory

2. 计算最优的 kernel 执行顺序

3. 预取 kernel 参数和 constant memory

2.3 Graph 更新机制

当模型参数或输入形状变化时，需要更新 Graph：

// 更新已实例化的 Graph
cudaGraphExecUpdateResultInfo updateResult;
cudaGraphExecUpdate(graphExec, newGraph, &updateResult);

// 或者使用 cudaGraphExecKernelNodeSetParams 更新单个节点参数
cudaKernelNodeParams newParams = {0};
newParams.func = (void*)my_kernel;
newParams.gridDim = dim3(128, 1, 1);
newParams.blockDim = dim3(256, 1, 1);
cudaGraphExecKernelNodeSetParams(graphExec, kernelNode, &newParams);

三、动态形状：CUDA Graph 的阿喀琉斯之踵

3.1 问题的本质

CUDA Graph 在实例化时固化了所有内存分配的大小。LLM 推理中的序列长度在运行时变化（beam search 的不同路径、batch 中不同请求的长度差异），这导致：

场景 A：Prefill 阶段——输入 prompt 长度从 1 到 4096 tokens 不等，attention 的计算量随长度平方增长。

场景 B：Decoding 阶段——每生成一个 token，KV cache 增长一层，memory footprint 持续变化（虽然单次增量小，但累积效应显著）。

场景 C：Batch 动态性——Continuous batching 下，新请求加入、已完成请求退出，batch size 随时变化。

3.2 朴素方案的局限

方案一：为每种形状独立录制 Graph

方案二：按最大形状录制，短序列浪费计算

当最长 prompt 为 4096 tokens 而大部分请求只有 128 tokens 时，GPU 利用率可能低至 30%。

方案三：每个 sequence 重新录制

录制开销（通常 5-50ms）在短序列场景下完全不可接受。

四、生产级方案：分段捕获 + 内存池 + 参数热更新

4.1 核心思想：Chunked Prefill + Fixed-Shape Decoding

现代引擎（vLLM、TensorRT-LLM）的通用做法是：

┌─────────────────────────────────────────────────────┐
│ Prefill 阶段 (variable length)                       │
│  ┌──────┐ ┌──────┐ ┌──────┐                        │
│  │chunk1│→│chunk2│→│chunk3│  (each chunk = 512 tok) │
│  └──────┘ └──────┘ └──────┘                        │
│  → 使用 CUDA Graph 捕获固定大小的 chunk              │
├─────────────────────────────────────────────────────┤
│ Decoding 阶段 (fixed shape per step)                 │
│  ┌───┐ ┌───┐ ┌───┐                                  │
│  │ d1│→│ d2│→│ d3│  (每个 decode step shape 固定)   │
│  └───┘ └───┘ └───┘                                  │
│  → Graph 只需捕获一次，复用全部生成过程              │
└─────────────────────────────────────────────────────┘

4.2 Chunked Prefill 的 CUDA Graph 实现

class ChunkedPrefillGraph {
public:
    static constexpr int CHUNK_SIZE = 512;
    static constexpr int MAX_HIDDEN = 8192;
    
    struct GraphBundle {
        cudaGraph_t graph;
        cudaGraphExec_t exec;
        int seq_len;  // 该 graph 对应的序列长度
    };

private:
    std::vector<GraphBundle> graphs_;
    int num_layers_;
    int hidden_dim_;

public:
    void capture(int num_layers, int hidden_dim) {
        num_layers_ = num_layers;
        hidden_dim_ = hidden_dim;
        
        // 为每种 chunk 大小录制 Graph
        for (int seq_len : {128, 256, 512, 1024, 2048, 4096}) {
            capture_for_length(seq_len);
        }
    }
    
    void capture_for_length(int seq_len) {
        cudaGraph_t graph;
        cudaStreamBeginCapture(stream_, cudaStreamCaptureModeGlobal);
        
        // 分配固定大小的缓冲区
        // Q: [batch, num_heads, seq_len, head_dim]
        // K: [batch, num_kv_heads, seq_len + past_kv_len, head_dim]  
        // 注意：这里 past_kv_len 固定为最大值
        void *q_buf, *k_buf, *v_buf, *out_buf;
        size_t q_size = batch_size_ * num_heads_ * seq_len * head_dim_ * sizeof(float16);
        cudaMalloc(&q_buf, q_size);
        // ... 分配其他 buffer
        
        for (int layer = 0; layer < num_layers_; layer++) {
            // Attention kernel with tiling
            flash_attention_kernel<<<grid, block, shm_size, stream_>>>(
                q_buf, k_buf, v_buf, out_buf,
                batch_size_, seq_len, num_heads_, head_dim_
            );
            
            // FFN with GEMM
            cublasGemmEx(handle_, ...);  // Graph-captured GEMM via cuBLAS LT
            
            // LayerNorm + Residual
            fused_layernorm_residual<<<grid, block, 0, stream_>>>(...);
        }
        
        cudaStreamEndCapture(stream_, &graph);
        
        cudaGraphExec_t exec;
        cudaGraphInstantiate(&exec, graph, NULL, NULL, 0);
        
        graphs_.push_back({graph, exec, seq_len});
    }
    
    void forward(int actual_seq_len, void* output) {
        // 找到能覆盖实际长度的最小 Graph
        GraphBundle* best = nullptr;
        for (auto& g : graphs_) {
            if (g.seq_len >= actual_seq_len) {
                best = &g;
                break;
            }
        }
        
        // 更新参数并启动
        update_graph_params(best, actual_seq_len);
        cudaGraphLaunch(best->exec, stream_);
    }
};

4.3 内存池设计：两阶段分配器

class CUDAGraphMemoryPool {
public:
    struct PoolConfig {
        size_t kv_cache_per_token;    // 每 token 的 KV cache 大小
        size_t activation_peak;       // 最大激活值大小
        size_t workspace_size;        // cuBLAS/cuDNN workspace
        int max_batch_size;
        int max_seq_len;
    };

private:
    // 阶段 1：Persistent buffers（录制一次，永不释放）
    struct PersistentPools {
        void* kv_cache;           // [max_batch, max_seq, num_layers, num_heads, head_dim]
        void* query_projections;  // 权重矩阵
        void* attention_output;
        void* ffn_intermediate;
    } persistent_;

    // 阶段 2：Reusable scratch space（在 Graph 之间复用）
    struct ScratchPool {
        void* buffer;
        size_t capacity;
        std::atomic<bool> in_use{false};
    };
    
    std::vector<ScratchPool> scratch_pools_;

public:
    cudaError_t allocate_persistent(const PoolConfig& config) {
        // 计算总需求并一次性分配
        size_t total_persistent = 
            config.kv_cache_per_token * config.max_batch_size * config.max_seq_len +
            config.activation_peak +
            config.workspace_size;
        
        // 使用 cudaMallocAsync 支持 pool 语义
        return cudaMalloc(&persistent_.kv_cache, total_persistent);
    }
    
    // 获取 scratch 空间（无锁快速路径）
    void* acquire_scratch(size_t min_size) {
        for (auto& pool : scratch_pools_) {
            if (!pool.in_use.load() && pool.capacity >= min_size) {
                bool expected = false;
                if (pool.in_use.compare_exchange_strong(expected, true)) {
                    return pool.buffer;
                }
            }
        }
        // 慢速路径：分配新 pool
        return create_new_pool(min_size);
    }
    
    void release_scratch(void* ptr) {
        for (auto& pool : scratch_pools_) {
            if (pool.buffer == ptr) {
                pool.in_use.store(false);
                return;
            }
        }
    }
};

4.4 参数热更新：避免重复录制的关键

class GraphExecutor {
    cudaGraphExec_t exec_;
    std::unordered_map<std::string, cudaGraphNode_t> named_nodes_;
    
public:
    // 更新 attention kernel 的序列长度参数
    void update_attention_params(int actual_seq_len, int past_kv_len) {
        auto it = named_nodes_.find("flash_attention");
        if (it == named_nodes_.end()) return;
        
        // Flash Attention kernel 参数结构
        struct FA_Params {
            int batch_size;
            int seq_len;
            int past_kv_len;
            int num_heads;
            int head_dim;
            half* Q;
            half* K;
            half* V;
            half* O;
        };
        
        FA_Params* params;
        cudaGraphKernelNodeGetParams(it->second, (cudaKernelNodeParams*)&params);
        
        // 只更新变化的部分，不改变图结构
        ((FA_Params*)params)->seq_len = actual_seq_len;
        ((FA_Params*)params)->past_kv_len = past_kv_len;
        
        cudaGraphExecKernelNodeSetParams(exec_, it->second, 
                                         (const cudaKernelNodeParams*)&params);
    }
    
    void execute() {
        cudaGraphLaunch(exec_, stream_);
        cudaStreamSynchronize(stream_);
    }
};

五、实战优化：Flash Attention + CUDA Graph 的协同设计

5.1 Flash Attention 的 tiling 策略

Flash Attention 将 attention 计算分 tile 处理，避免实例化完整的 [seq_len, seq_len] 注意力矩阵：

Normal Attention:
Q [1, 4096, 8192] × K^T [1, 8192, 4096] = S [1, 4096, 4096]  ← 64GB 中间矩阵
S × V = O [1, 4096, 8192]

Flash Attention (Br=256, Bc=256, Tr=16, Tc=16):
└─ 外循环: K, V 分块加载到 SRAM
└─ 内循环: Q 分块计算，online softmax 增量合并
└─ 峰值显存: O(Br × d + Br × Bc) ≈ 数 MB

5.2 CUDA Graph 友好的 Flash Attention 封装

// 关键：使用固定 tile size，让 CUDA Graph 可以安全捕获
void flash_attn_packed_fwd(
    half* __restrict__ q,      // [total_q, num_heads, head_dim]
    half* __restrict__ k,      // [total_kv, num_kv_heads, head_dim]  
    half* __restrict__ v,
    half* __restrict__ out,
    int* __restrict__ cu_seqlens_q,  // cumulative sequence lengths
    int* __restrict__ cu_seqlens_k,
    int max_seqlen_q,
    int max_seqlen_k,
    int num_heads,
    int num_kv_heads,
    int head_dim,
    cudaStream_t stream
) {
    // 使用 fixed tile size = 128，确保每次 kernel launch configuration 一致
    constexpr int kBlockM = 128;
    constexpr int kBlockN = 128;
    
    int const num_m_block = (max_seqlen_q + kBlockM - 1) / kBlockM;
    int const num_n_block = (max_seqlen_k + kBlockN - 1) / kBlockN;
    
    // 通过 max_seqlen 计算 grid，确保 graph 捕获时 grid 维度固定
    dim3 grid(num_m_block, num_heads, batch_size);
    dim3 block(kBlockM, 1);  // 256 threads per block
    
    flash_attention_kernel<kBlockM, kBlockN, kHeadDim>
        <<<grid, block, 0, stream>>>(
            q, k, v, out,
            cu_seqlens_q, cu_seqlens_k,
            max_seqlen_q, max_seqlen_k
        );
}

5.3 实际性能对比（A100-80GB, LLaMA-2-70B, INT4 量化）

观察：

• CUDA Graph 在 decode 阶段收益最显著（+67%），因为前向传播时间最短

• Chunked Prefill 相比 Static Shape 性能损失仅 2%，但显存节省 8%

• 多次录制可以通过 Graph Instantiation 的 lazy 执行摊销

六、高级技巧与避坑指南

6.1 避免 Graph 捕获期间的内存分配

// ❌ 错误：捕获期间分配内存会导致 Graph 记录 cudaMalloc 节点
cudaStreamBeginCapture(stream);
cudaMalloc(&temp_buffer, size);  // 每次回放都会执行分配！
my_kernel<<<..., stream>>>(temp_buffer);
cudaStreamEndCapture(stream, &graph);

// ✅ 正确：使用预分配的 buffer
void* temp_buffer;  // 预分配
cudaMalloc(&temp_buffer, max_size);

cudaStreamBeginCapture(stream);
my_kernel<<<..., stream>>>(temp_buffer);  // 录制时引用已知地址
cudaStreamEndCapture(stream, &graph);

6.2 使用 CUDA Events 进行跨 Graph 同步

class MultiGraphPipeline {
    cudaGraphExec_t prefill_graph_;
    cudaGraphExec_t decode_graph_;
    cudaEvent_t prefill_done_;
    
public:
    void execute_mixed_batch(RequestBatch& batch) {
        // Prefill 长序列
        for (auto& req : batch.prefills) {
            cudaGraphLaunch(prefill_graph_, compute_stream_);
        }
        
        // 通知 decode stream
        cudaEventRecord(prefill_done_, compute_stream_);
        cudaStreamWaitEvent(decode_stream_, prefill_done_);
        
        // 执行 decode（与 prefill 并行）
        cudaGraphLaunch(decode_graph_, decode_stream_);
    }
};

6.3 Graph 池化模式（Graph Pool Pattern）

class GraphPool {
    std::array<cudaGraphExec_t, 32> pool_;  // 预分配 N 个 Graph
    std::array<bool, 32> available_;
    
public:
    cudaGraphExec_t acquire() {
        for (int i = 0; i < 32; i++) {
            bool expected = true;
            if (available_[i].compare_exchange_strong(expected, false)) {
                return pool_[i];
            }
        }
        // 降级：使用传统执行模式
        return nullptr;
    }
    
    void release(cudaGraphExec_t graph) {
        for (int i = 0; i < 32; i++) {
            if (pool_[i] == graph) {
                available_[i].store(true);
                return;
            }
        }
    }
};

七、展望：CUDA Graph 与下一代硬件的协同演进

NVIDIA 的 Hopper 架构引入了 Thread Block Cluster 和 Distributed Shared Memory，为 CUDA Graph 带来了新的优化空间：

1. Cluster-aware Graph：跨 SM 的异步数据交换可以通过 Graph 节点捕获，减少 barrier overhead。

2. Graph Memory Node（CUDA 12+）：允许在 Graph 中包含 cudaMallocAsync-style 的动态分配，从根本上缓解动态形状限制。

3. Conditional Nodes（CUDA 12.3+）：支持在 Graph 内部做条件分支，使得 batch size 变化只需切换 Graph 分支而非重新录制。

TensorRT-LLM 和 vLLM 正在积极采纳这些特性。未来，结合 CUDA Graph 与 CUDA Dynamic Parallelism，LLM 推理引擎有望实现真正的"任意形状、零开销"执行。

总结

CUDA Graph 在 LLM 推理中的应用，本质上是用空间换时间的经典权衡：

• 牺牲一定的灵活性（静态形状约束），换取 launch overhead 的消除

• 通过 Chunked Prefill + Fixed-Shape Decoding 平衡两者

• 内存池设计将显存碎片化降至最低

• 参数热更新机制使得单次录制可服务多种输入

理解并善用这些技术，是在 GPU 算力日趋紧张的背景下，榨取每一点推理性能的关键所在。

核心要点回顾：

• CUDA Graph 将 O(N) launch overhead 降至 O(1)

• 动态形状是核心挑战，需要 Chunked Prefill 策略

• 两阶段内存池（Persistent + Scratch）解决显存管理

• 参数热更新避免重复录制，是生产环境关键能力

• Hopper 的 Graph Memory Node 和 Conditional Nodes 预示更灵活的 Graph 未来