从零实现 LLM Inference:043. Streaming Benchmark 开关(paged-attn / CUDA Graph / NVTX)
streaming 场景想看 TTFT/TPOT/ITL 的 p99,但之前不好一条命令切换 decode 的快路径,也不好给 nsys 挂 NVTX。补齐 benchmark 的三个开关。
streaming 场景想看 TTFT/TPOT/ITL 的 p99,但之前不好一条命令切换 decode 的快路径,也不好给 nsys 挂 NVTX。补齐 benchmark 的三个开关。
decode 每层都要把 (k,v) 写进 KV cache。原来是 index_select + index_put 四连。先用一个 Triton kernel fuse 成一次写入,并加上 batch size gate,避免小 batch 退化。
decode 已经很快了但还不够?很多时候慢在 sampling:每步 B 次 .item() 会把 GPU pipeline 打散。把采样做成 batch,一步只同步一次。
decode 热路径里 kernel launch 太密?把一次 decode step 捕获成 CUDA Graph,replay 省掉大量 CPU dispatch。
把 paged-attn decode kernel 的 num_warps/num_stages 从“拍脑袋常量”变成 Triton autotune;结果在 decode 热路径里反而回退了。