PyTorch自定义LLM推理中CUDA Graph生成显存溢出问题排查

核心问题分析 #

你当前的实现存在两个致命的显存浪费点：

1. 每个CUDA Graph硬编码了graph_steps次完整的token生成流程，相当于把graph_steps个单步操作打包进一个Graph。每个这样的Graph会完整记录所有中间操作的依赖，显存占用直接是单步生成Graph的graph_steps倍。
2. 你创建了max_seq_len//graph_steps个这样的Graph，总显存占用就变成了max_seq_len × 单步Graph大小，这直接撑爆了A100的80GB显存。

至于你怀疑的切片导致全局张量复制：CUDA Graph不会复制全局张量本身（切片是视图，不占额外显存），但多次重复的操作记录才是显存爆炸的根本原因。

修正方案 #

放弃为每个分段创建独立Graph的思路，改为捕获单步token生成的Graph，通过动态更新索引来重复复用这个Graph，具体步骤如下：

1. 预先分配全局张量

确保所有需要持久化的张量（KV缓存、生成的token序列）都是预先分配好的全局张量，避免动态分配：

# 示例：预分配KV缓存（根据你的模型维度调整）
global_kv_cache = torch.zeros(
    (num_layers, 2, batch_size, head_dim, max_seq_len),
    dtype=torch.float16,
    device="cuda"
)
global_tokens = torch.zeros(max_seq_len, dtype=torch.long, device="cuda")

2. 捕获单步生成的CUDA Graph

使用可修改的张量存储当前位置索引，确保Graph可以动态适配不同的KV缓存位置：

# 用可修改的张量存储当前生成位置（作为Graph的输入占位符）
current_pos = torch.tensor(0, dtype=torch.long, device="cuda")

single_step_graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(single_step_graph):
    # 1. 基于current_pos访问对应位置的KV缓存
    current_k = global_kv_cache[:, 0, :, :, current_pos]
    current_v = global_kv_cache[:, 1, :, :, current_pos]
    
    # 2. 执行LLM前向推理（替换为你的模型逻辑：Attention、MLP等）
    new_token, new_k, new_v = llm_forward(global_tokens[:current_pos+1], current_k, current_v)
    
    # 3. 原地更新全局张量（关键：必须用原地操作，避免创建新张量）
    global_kv_cache[:, 0, :, :, current_pos+1] = new_k
    global_kv_cache[:, 1, :, :, current_pos+1] = new_v
    global_tokens[current_pos+1] = new_token
    
    # 4. 原地更新当前位置（或外部手动更新）
    current_pos.add_(1)

3. 推理阶段复用Graph

通过更新current_pos的值，重复执行同一个Graph即可完成所有token的生成：

# 初始化位置
current_pos.fill_(0)
# 执行max_seq_len次生成（或直到终止条件）
for _ in range(max_seq_len - 1):
    single_step_graph.replay()

额外优化建议 #

• 如果需要批量执行多个步骤（比如一次执行graph_steps次），可以在外部循环中调用replay()多次，不需要重新捕获Graph。
• 确保模型的所有前向操作都是静态图兼容的，避免动态分支、动态张量形状（你已经用预分配全局张量做到了这一点）。
• 可以使用torch.cuda.make_graphed_callable包装单步生成函数，进一步简化Graph的复用逻辑，自动处理输入参数的更新。

内容的提问来源于stack exchange，提问作者PatchyTIS