【模型推理】浅谈CUDA Graph(WIP)

// Stream 模型：线性队列，CPU 逐条提交
cudaMemcpyAsync(d_in, h_in, size, H2D, stream);
kernelA<<<grid, block, 0, stream>>>(d_in, d_tmp);  // A
kernelB<<<grid, block, 0, stream>>>(d_tmp, d_out); // B，必须等 A 完事

cudaStream_t s1, s2;
cudaEvent_t ev;
// ... 创建 stream 和 event

kernelA<<<grid, block, 0, s1>>>();
cudaEventRecord(ev, s1);           // A 完了发个信号

kernelB<<<grid, block, 0, s2>>>();
cudaStreamWaitEvent(s2, ev);       // B 等 A
kernelC<<<grid, block, 0, s2>>>(); // C 也等 A，但 B 和 C 可以并行

cudaGraph_t graph;
cudaGraphCreate(&graph);

// 添加一个 kernel 节点
cudaGraphNode_t kernelNode;
cudaKernelNodeParams params = {...};
cudaGraphAddKernelNode(&kernelNode, graph, NULL, 0, &params);

// 添加依赖边
cudaGraphAddDependencies(graph, &nodeA, &nodeB, 1);  // A → B

cudaGraph_t graph;
cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal);

// 这里写你原来的代码，会被自动录下来
kernelA<<<grid, block, 0, stream>>>();
kernelB<<<grid, block, 0, stream>>>();
// ...

cudaStreamEndCapture(stream, &graph);

import torch

# Enable CUDA Graph mode in PyTorch
model = YourModel().cuda()
static_input = torch.randn(32, 3, 224, 224, device='cuda')

# Warmup
for _ in range(3):
    _ = model(static_input)

# Capture the graph
g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g):
    static_output = model(static_input)

# Training loop - replay the graph
for data in dataloader:
    static_input.copy_(data)  # Update input in-place
    g.replay()                # Execute captured operations

# 1. 预分配固定 buffer（整个生命周期不释放）
static_input = torch.empty(max_batch, max_seq_len, device='cuda')

# 2. 每次推理：往固定地址填新数据
static_input.copy_(new_input)

# 3. replay：图从同一个地址读
g.replay()

# 4. 读结果：从固定地址拷出来
result = static_output.clone()

# 不安全
if x.sum() > 0:
    y = x * 2
else:
    y = x + 1

# 安全 —— 用 mask 代替分支
mask = (x > 0).float()
y = mask * (x * 2) + (1 - mask) * (x + 1)

# 不安全
for i in range(batch_size):  # batch_size 是 tensor 的值
    process(x[i])

# 安全 —— 用固定次数的循环
for i in range(MAX_BATCH):  # MAX_BATCH 是常量
    if i < batch_size:
        process(x[i])

# 错误：capture 期间 sync 会卡住
with torch.cuda.graph(g):
    x = model(input)
    torch.cuda.synchronize()  # 死锁！操作根本没执行，你在等一个永远不会完成的事

# 错误：capture 期间 tensor 的值不存在
with torch.cuda.graph(g):
    x = model(input)
    print(x[0])       # 报错：无法读取值
    if x.sum() > 0:   # 报错：条件依赖值
        ...

# 错误：依赖 tensor 值的控制流
with torch.cuda.graph(g):
    for i in range(seq_len):  # seq_len 如果是 tensor 的值，不行
        process(x[i])

# 正确：用固定次数 + mask
with torch.cuda.graph(g):
    for i in range(MAX_LEN):  # MAX_LEN 是 Python 常量
        mask = (i < seq_len).float()
        out += mask * process(x[i])

from __future__ import annotations

import json
import statistics
import time

import torch
from transformers import StaticCache

MODEL_ID = "Qwen/Qwen3-0.6B"
DEVICE = "cuda:0"
DTYPE = torch.bfloat16

# 测量参数
PROMPT = "The key difference between CUDA streams and CUDA graphs is"
NUM_DECODE_TOKENS = 128
WARMUP_STEPS = 10

RESULTS_DIR = Path(__file__).resolve().parent.parent / "results"


@dataclass
class LoadedModel:
    model: torch.nn.Module
    tokenizer: AutoTokenizer


def load_model() -> LoadedModel:
    """加载 Qwen3-0.6B。eval + no_grad 在调用侧处理。"""
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        torch_dtype=DTYPE,
        attn_implementation="sdpa",
    ).to(DEVICE)
    model.eval()
    return LoadedModel(model=model, tokenizer=tokenizer)


def encode_prompt(tokenizer, prompt: str = PROMPT) -> torch.Tensor:
    """Prompt → input_ids on device. Shape: [1, prompt_len]."""
    return tokenizer(prompt, return_tensors="pt").input_ids.to(DEVICE)


def _prefill(model, input_ids, cache):
    """Prefill 走 eager. 返回首个 decode token 和 next cache_position."""
    prompt_len = input_ids.shape[1]
    cache_position = torch.arange(prompt_len, dtype=torch.int64, device=DEVICE)
    out = model(
        input_ids=input_ids,
        past_key_values=cache,
        cache_position=cache_position,
        use_cache=True,
    )
    next_tokens = out.logits[:, -1:, :].argmax(dim=-1)
    next_cache_position = cache_position[-1:] + 1
    return next_tokens, next_cache_position


def _eager_decode_step(model, cache, input_ids, cache_position):
    """跑一步 decode. 用于 warmup / 为 capture 准备 cuBLAS workspace."""
    out = model(
        input_ids=input_ids,
        past_key_values=cache,
        cache_position=cache_position,
        use_cache=True,
    )
    return out


@torch.no_grad()
def run_graphed() -> dict:
    loaded = load_model()
    model, tokenizer = loaded.model, loaded.tokenizer
    input_ids = encode_prompt(tokenizer)
    prompt_len = input_ids.shape[1]
    max_cache_len = prompt_len + NUM_DECODE_TOKENS + 16

    # ─────────────────────────────────────────────
    # 一次性 cache: 从头到尾只建一个
    # 原因: graph 捕获到的 cache 指针必须在 replay 时依然有效
    # ─────────────────────────────────────────────
    cache = StaticCache(
        config=model.config, max_cache_len=max_cache_len, device=DEVICE, dtype=DTYPE,
    )

    print(f"[setup] prefill + {WARMUP_STEPS} eager warmup decode steps ...")
    next_tokens, next_cache_position = _prefill(model, input_ids, cache)
    out = None
    for _ in range(WARMUP_STEPS):
        out = _eager_decode_step(model, cache, next_tokens, next_cache_position)
        next_tokens = out.logits[:, -1:, :].argmax(dim=-1)
        next_cache_position = next_cache_position + 1
    torch.cuda.synchronize()

	# 静态输入输出
    static_input_ids = next_tokens.clone()                # [1, 1]
    static_cache_position = next_cache_position.clone()   # [1]

    print("[capture] recording decode graph ...")
    
    s = torch.cuda.Stream()
    g = torch.cuda.CUDAGraph()
    
    s.wait_stream(torch.cuda.current_stream())
    with torch.cuda.stream(s):
        for _ in range(3):
            out = model(
                input_ids=static_input_ids,
                past_key_values=cache,
                cache_position=static_cache_position,
                use_cache=True,
            )
    torch.cuda.current_stream().wait_stream(s)
    
    
    with torch.cuda.graph(g):
        out = model(
                input_ids=static_input_ids,
                past_key_values=cache,
                cache_position=static_cache_position,
                use_cache=True,
            )
    static_output_logits = out.logits 

    # 重置 cache 状态 + 重做 prefill
    # 关键: StaticCache 内部有一个 cumulative_length 计数器, 每次 forward 都 +=
    # input_len (不管 cache_position 是否复用). 如果不 reset, 它会一路涨到超过
    # max_cache_len, 导致某处 index_copy_ 触发 "index out of bounds".
    # cache.reset() 会把 keys/values 清零, 并把 cumulative_length.zero_().
    # 注意: 这只是零内容, 地址不变——这正是 graph 依赖的前提, 所以 reset 不影响 g.
    cache.reset()

    next_tokens, next_cache_position = _prefill(model, input_ids, cache)
    static_input_ids.copy_(next_tokens)
    static_cache_position.copy_(next_cache_position)
    torch.cuda.synchronize()


    # 测量: 用 graph.replay() 代替 model() 做 decode loop
    events_pair: list[tuple[torch.cuda.Event, torch.cuda.Event]] = []

    t0 = time.perf_counter()
    for step in range(NUM_DECODE_TOKENS - 1):  # -1 因为 prefill 已经产生了第一个 token
        step_start = torch.cuda.Event(enable_timing=True)
        step_end = torch.cuda.Event(enable_timing=True)

        step_start.record()
        g.replay()
        step_end.record()
        events_pair.append((step_start, step_end))

        # 用 replay 的输出产生下一步输入. copy_() 原地更新静态 buffer
        next_tokens = static_output_logits[:, -1:, :].argmax(dim=-1)
        static_input_ids.copy_(next_tokens)
        static_cache_position.add_(1)  # 原地 +1

    torch.cuda.synchronize()
    t1 = time.perf_counter()
    total_wall_ms = (t1 - t0) * 1000
    per_step_gpu_ms = [s.elapsed_time(e) for s, e in events_pair]

	
	# 统计
    decode_ms = per_step_gpu_ms  # 这里全部是 decode (prefill 在外面)
    result = {
        "model": MODEL_ID,
        "dtype": "bfloat16",
        "cache": "static",
        "attention": "sdpa",
        "mode": "graphed",
        "max_cache_len": max_cache_len,
        "prompt_len": prompt_len,
        "num_decode_tokens": len(decode_ms),
        "decode_per_step_ms": decode_ms,
        "decode_stats": {
            "mean": statistics.mean(decode_ms),
            "median": statistics.median(decode_ms),
            "stdev": statistics.stdev(decode_ms) if len(decode_ms) > 1 else 0.0,
            "min": min(decode_ms),
            "max": max(decode_ms),
        },
        "total_wall_ms": total_wall_ms,
        "total_gpu_ms": sum(per_step_gpu_ms),
        "cpu_overhead_ms": total_wall_ms - sum(per_step_gpu_ms),
    }


    print(f"\n[graphed] decode mean:  {result['decode_stats']['mean']:.3f} ms/step")
    print(f"[graphed] decode stdev: {result['decode_stats']['stdev']:.3f} ms")
    print(f"[graphed] total wall:   {total_wall_ms:.1f} ms")
    print(f"[graphed] CPU overhead: {result['cpu_overhead_ms']:.1f} ms "
          f"({result['cpu_overhead_ms']/total_wall_ms*100:.2f}%)")
    print(f"[graphed] saved → {out_path}")
    return result


if __name__ == "__main__":
    run_graphed()