万字长文!C++/CUDA 视角硬核解析 TensorRT-LLM XQA 算子:从 Roofline 拐点到动态 Split-K 实现
在大模型线上推理服务压测中,经常出现显存占用与 GPU 利用率(GPU-Util)双双报满,但实际 Token 生成吞吐(TPS)远落后于预期、单步生成延迟(per-token latency)居高不下的情况。这并非系统死锁或主机端锁竞争所致,而是自回归解码(Decode)阶段固有的体系结构特性使然:计算模式从预填充(Prefill)阶段的高密度通用矩阵乘法(GEMM)退化为显存带宽受限的向量-矩阵乘法(GEMV)。算术强度(Operational Intensity)相较 Prefill 阶段下降了两个数量级,计算核心(Tensor Core)长期处于等待片外显存加载的停顿状态。同时,在小 Batch 或单会话场景下,传统网格划分产生的线程块数量过少,无法覆盖现代 GPU 上的全部流式多处理器(SM),导致硬件出现明显的并行度饥饿。针对带宽利用率不足与硬件空闲的问题,TensorRT-LLM 引入了 XQA(eXtended Query Attention)算子,结合 GQA 共享内存广播、Split-K 跨 Block 在线 Softmax 状态重整化以及 Paged KV Cache 二级寻址机制,从硬件微架构层面重构了 Decode 阶段的执行管线。本文将从硬件体系结构、数学推导与工程实现三个维度,系统解析该方案的技术细节。+----------------------------------------------------+ | 自回归生成阶段硬件瓶颈全景分析 | | | | Prefill