0x0. 前言

今天聊一个最近有趣的发现，那就是模型推理时是否应该在 GTX 4090 上开启 cuda graph ？在 GTX 4090 上用推理框架如VLLM/SGLang等，什么情况下才应该开启 CUDA Graph？目前只能说一下我的观察过程和结论，背后可能的原因也请大佬不吝赐教。

0x1. 问题发生的背景

某天，我想看一下在 GTX 4090 单卡情况下使用VLLM和Qwen2-7B时离线推理一个 prompt 的时候相比于 HuggingFace 原始的推理有多大的性能提升。

这里主要关注decoding过程中每个iter的速度，因为prefill只有一次，且 VLLM/SGLang 都不会通过 cuda-graph 来加速prefill过程，并且decoding会触发频繁的 cuda kernel launch。

然后，我写了下面2个脚本，分别用于测试VLLM和HuggingFace Qwen2-7B的推理性能，我使用nsight system来profile，脚本开头是profile的指令。

vllm 推理脚本

# /opt/nvidia/nsight-systems/2024.5.1/bin/nsys profile --trace-fork-before-exec=true --cuda-graph-trace=node -o vllm_qwen2.5_7b_eager python3 debug.py
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

import nvtx
import torch
from vllm import LLM, SamplingParams

# Sample prompts.
prompts = "帮我计划一次去北京的旅行，我想明年春天出发，大概五天的行程。"
# Create a sampling params object.
sampling_params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=512)

# Create an LLM.
llm = LLM(model="/mnt/bbuf/Qwen2.5-7B-Instruct/", enforce_eager=True)
# Generate texts from the prompts. The output is a list of RequestOutput objects
# that contain the prompt, generated text, and other information.
# warmup
for _ in range(2):
    outputs = llm.generate(prompts, sampling_params)

torch.cuda.synchronize()

# profile
for i in range(20):
    with nvtx.annotate(f"step={i}", color="blue"):
        outputs = llm.generate(prompts, sampling_params)

# Print the outputs.
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

注意，这个脚本中我暂时开启了 enforce_eager=True 来关闭 CUDA Graph。

HuggingFace 推理脚本

# /opt/nvidia/nsight-systems/2024.5.1/bin/nsys profile --trace-fork-before-exec=true --cuda-graph-trace=node -o hf_qwen2.5_7b_flash_attn python3 debug.py
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

import nvtx
import torch

from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "/mnt/bbuf/Qwen2.5-7B-Instruct"

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

prompt = "帮我计划一次去北京的旅行，我想明年春天出发，大概五天的行程。"

model_inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

# warmup
for _ in range(2):
    generated_ids = model.generate(
        **model_inputs,
        max_new_tokens=512
    )

torch.cuda.synchronize()
# profile

for i in range(20):
    with nvtx.annotate(f"step={i}", color="blue"):
        
        generated_ids = model.generate(
            **model_inputs,
            max_new_tokens=512
        )

generated_ids = [
    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
]

response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
print(response)

nsys结果分析

• vllm

• hf

都使用 Eager 推理时，我发现 VLLM 的一个decoding的iter 15.8ms，然后 HF 的一个decoding的iter 18.1ms。关注到decoding阶段kernel launch的速度都非常快，ns级别，这种情况CUDA Graph应该无法发挥出作用。至于15.8ms和18.1ms的差异，来源在于fused rope，fused rmsnorm，packed qkv linear，我把这几个组件调整成一样HF就可以和VLLM具有相同的单卡推理性能。

验证一下，我把上面VLLM 脚本里的 enforce_eager=True 去掉，开启 CUDA Graph，再跑一遍，nsys结果如下：

decoding一个iter的时间和 Eager 模式是一样的。

现在引出了本文的问题，什么时候在 GTX 4090 上开启 CUDA Graph？

相比之下，如果在A800上执行上面的脚本，如果不开启cuda graph则一个decoding的iter需要37ms，开启之后只需要13ms，差异非常明显。

0x2. SGLang推理时 CUDA Graph 开启的观察

为了探索在 GTX 4090 推理模型时什么情况下需要打开 CUDA Graph，我基于 SGLang 做了一系列的实验。

我基于 SGlang v0.3.6，使用sharegpt的数据来测试了以下模型：