TensorRT-LLM 源码完全拆解（二十九）：XQA 解码注意力内核——分页 KV、GQA 与 FP8 如何被 Warp 特化流水吃下

为什么单独拿 XQA 出来讲

前面几篇把 TensorRT-LLM 的注意力后端从「选谁」讲到「元数据怎么准备」：AttentionMetadata.prepare() 把块句柄翻译成各后端需要的索引数组，FlashInfer 负责分页 KV 的通用路径，TRTLLM-Gen 的 FMHA 接管 Blackwell 上的新一代内核。可真正撑起 Hopper 及以前各代 GPU 上 decode 注意力的，是一个长期藏在 cpp/kernels/xqa/ 目录里的独立内核族——XQA（官方 README 的定位是 “a set of optimized kernels for generation-phase MQA/GQA”）。

XQA 值得单独拆，原因有三点：

1. 它把 decode 注意力的三个约束——每次只算一个 query token、KV 在分页池里不连续、GQA 的多个 Q 头共享一份 KV——同时压进一个内核，而不是像训练侧那样假设 Q 是一个长矩阵。
2. 它没有走「一个模板打天下」的路，而是把 head_size、GQA 组大小、页大小、KV 精度、beam width、投机解码开关全部变成编译期宏，用 NVRTC 在运行时编译出只服务当前模型形状的 cubin。
3. 它把 Ampere 与 Hopper 分成两套完全不同的实现：Ampere 路径用 mma.sync 加软件流水，Hopper 路径用 TMA 加 warpgroup MMA，两者共享同一套 host 侧调度接口。

这一篇按「host 怎么决定用 XQA → 内核参数怎么变成宏 → Ampere 内核的 warp 分工 → Hopper 内核的三角色流水 → 分页 KV 与 FP8 的细节 → 多块模式与工程取舍」的顺序走。读完后应该能回答一个具体问题：一条 decode 请求进来，XQA 凭什么比通用 MMHA 快，又凭什么在某些形状下反而不如它。

先把 decode 注意力的约束写清楚

一个 decode step 里，每条请求只有 1 个新 query token（开启投机解码时是一小段草稿 token；HMMA 路径把「草稿长度 × GQA 组大小」压到 16 或 32 行，QGMMA 路径固定为 64 行）。设 batch 为 $B$、每个 KV 头的 query 头数为 $G$（GQA 的组大小）、KV 头数为 $H_{kv}$、每头维度为 $D_h$、上下文长度为 $L$。这一步要算的是：

其中 $K_{b,h}, V_{b,h}$ 的形状是 $L \times D_h$。三个数字决定了内核该怎么写：

• 算术强度：一次 decode 注意力对每个 KV 头要读 $2 \times L \times D_h$ 个元素，做 $2 \times G \times L \times D_h$ 次乘加。当 $G = 8$ 时，每读一个 KV 元素就能复用 8 次，算术强度是「读一次算 8 次」量级——比单头 MHA（$G=1$）好一个数量级，但离 GEMM 主循环那种「读一次算几十次」还差得远。所以 decode 注意力天然是访存受限的，内核设计的核心是「把 KV 流式搬进 SMEM 的同时，让 MMA 单元不闲着」。
• 并行度：一个 CTA 处理一个 KV 头时，$M$ 方向只有 $G \times \text{beam}$ 行（普通解码、$G=8$、beam=1 时就是 8 行）。以 Llama-70B 这类 64 Q 头 / 8 KV 头的模型为例，TP8 部署后每张卡只剩 1 个 KV 头和 8 个 Q 头；batch 开到 32，每卡也只有 32 个 CTA，在 132 个 SM 的 H100 上连一个 wave 的零头都填不满，而每个 CTA 内部又只有 8 行——两个方向都没有足够的并行度。
• KV 不连续：分页 KV 缓存里，一条序列的 $L$ 个 token 分散在 $\lceil L / P \rceil$ 个页里，每页 $P$ 个 token（$P$ 取 8/16/32/64/128）。内核不能假设 $K_{b,h}$ 是一块连续内存。

这三点共同决定了 XQA 的形态：窄的 $M$ 方向要靠 GQA 组内共享 KV 换取复用；并行度不足要靠「沿序列维再切一刀」（多块模式）补上；分页 KV 要靠「先算页号、再过 TMA 或手写 gather」绕开不连续性。

图里是 MHA、MQA 与 GQA 三种头共享方式：MHA 每个 Q 头配一份 KV，MQA 所有 Q 头共用一份 KV，GQA 把 Q 头分组、每组共用一份 KV。XQA 的 HEAD_GRP_SIZE 就是图中 GQA 的组大小 $G$；它决定了每个 KV 头要被多少个 Q 头复用，也决定了内核 $M$ 方向的行数。来源：GQA 论文（arXiv:2305.13245）。

通用 MMHA（masked multi-head attention）内核同样要处理这些约束，但它的做法更保守：K 与 V 被逐 token 搬进共享内存，线程按「每个线程负责 K 向量的一段、再跨 warp 归约」的方式算点积，Q 与 K 的相乘是 dot 而不是 tensor core 的 mma。它对任意头数都成立，但每个 token 的 KV 只被用「一次点积」的力度消费，没有为 GQA 组织起「一份 KV 喂 $G$ 个 Q 头」的 tensor core 瓦片，也没有为 decode 特有的窄 $M$ 设计过布局。XQA 则是从「GQA 解码」这个特定形状出发重新设计的。

GQA 的组内复用：XQA 的数据布局前提

GQA 在数学上是「多个 Q 头共享一份 K/V」，但在内核里它首先是一个 内存布局约定：同一个 KV 头对应的 $G$ 个 Q 头必须在内存里连续排列，且顺序与 KV 头一一对应。否则一次 TMA 或 cp.async 搬进来的 K 瓦片，就无法在同一个 MMA 里同时服务这 $G$ 个 Q 头。

XQA 把这件事写成了两个常量：HEAD_GRP_SIZE（即 $G$）和 HEAD_ELEMS（即 $D_h$）。内核里的行号直接就是「Q 头在组内的偏移」：

1
constexpr uint32_t nbValidRows = headGrpSize * beamWidth;
2
constexpr uint2 warpTile = {64, roundUp(nbValidRows, 16U)};

nbValidRows 是 $G \times \text{beam}$。普通解码、$G=8$、beam=1 时是 8 行，roundUp(8, 16) 得到 16 行的瓦片，其中后 8 行是填充位。填充位不是浪费——mma.sync 的 $M$ 粒度是 16，8 行本来就必须补到 16 才能发射指令。真正的浪费只发生在「补出来的 8 行」上，而 K/V 的读取量完全由 $D_h$ 与 token 数决定，不受 $M$ 填充影响。

组内复用的收益可以精确算出来。一个 K 瓦片是 $64 \times D_h$，搬进 SMEM 要做 $64 D_h$ 次加载；同一个瓦片被 $G$ 个 Q 头共用，做 $G \times 64 \times D_h$ 次乘加。于是每个 KV 元素的复用次数是 $G$，GQA 把每个 KV 元素的复用次数从 MHA 的 1 提到了 8（$G=8$ 时），这是 decode 访存效率提升的直接来源。这就是为什么 XQA 的准入条件里专门把「原版 MHA 且 beam=1」排除掉——$G=1$ 时这份复用不存在，XQA 为组内布局付出的间接寻址反而成了纯开销。

输出侧也遵循同一约定。更完整的 GQA 机制与推理收益可参考 GQA 完全拆解：分组查询注意力，用 1/8 的 KV 缓存保住 MHA 的质量。写回时的偏移按「头组」而不是「头」计算：

1
uint32_t const outOffset = headGrpSize * (nbKHeads * (beamWidth * ctaInputTokBeg) + idxHeadGrp);

idxHeadGrp 是 KV 头编号，乘上 headGrpSize 就得到该组第一个 Q 头的全局位置；组内第 $g$ 个头再偏移 $g$。这也解释了为什么 XQA 要求 num_q_heads % num_kv_heads == 0——只有整除时，每个 KV 头才能对齐到一组固定大小的 Q 头，页表与头号才能用同一套乘法寻址。

XQA 在 TensorRT-LLM 里的位置：谁决定用它

XQA 并不是一个用户直接调用的 API。它挂在 TensorRT-LLM 的 decoder 注意力路径里，由 DecoderXQARunner 决定「这次前向到底用不用 XQA」，再用 CompileEngine 产出 cubin、用 run() 组装参数并发射。官方文档里对这一层的描述很直接：这是一个面向 MQA/GQA 生成阶段的优化，默认开启，但有一个启发式来决定用 XQA 还是 masked MHA，想强制使用可以设置 TRTLLM_FORCE_XQA=1。

决定流程的第一道关是 DecoderXQARunner::supportConfig()，它把三类内核的准入条件并起来：

1
bool DecoderXQARunner::supportConfig(XQAParams const& xqaParams, bool forConfigurePlugin) const
2
{
3
    return jit::supportConfigQGMMA(xqaParams, mSM, forConfigurePlugin)
4
        || jit::supportConfigHMMA(xqaParams, mSM, forConfigurePlugin)
5
        || jit::supportConfigMLA(xqaParams, mSM, forConfigurePlugin);
6
}

三类内核对应三种硬件画像：HMMA 是 mma.sync 版本，SM80 及以上都能跑；QGMMA 是 Hopper 专属的 TMA + warpgroup MMA 版本，只在 SM90 上；MLA 是 SM120 家族上专为 DeepSeek 式低秩 KV 写的版本。三者都要先过 supportConfigCommon()，这一层过滤掉的是「解码注意力语义不对」的情况：

其中「纯 MHA 不用 XQA」这一条容易被忽略。MHA 每条请求的 $G=1$，XQA 为 GQA 组织的「同一份 KV 喂给 $G$ 个 Q 头」的布局在这里没有收益，反而多了一层间接寻址。这一类请求交给 masked MHA 更合适。

过了准入之后还有第二道关 mayHavePerfGain()：即使配置支持，也要估计这颗 kernel 在当前 batch 下能不能填满 GPU。

1
int multiBlockCount = 1;
2
if (xqaParams.multi_block_mode)
3
{
4
    int historyLength = xqaParams.max_past_kv_length;
5
    multiBlockCount = std::max(1, historyLength / kMinHistoryTokensPerBlock);
6
}
7
int blockCount = numKVHeads * batchSize * multiBlockCount;
8
return static_cast<float>(blockCount) * kEnableMinBlockFactor >= static_cast<float>(mMultiProcessorCount);

kMinHistoryTokensPerBlock = 128、kEnableMinBlockFactor = 4.0。含义是：把每条序列按每 128 个历史 token 粗估出工作块数量；只有当这个数量的 4 倍还不少于 SM 数量时，才认为 XQA 值得上。换句话说，当上下文很短、batch 很小时，XQA 可能连一个 SM 的零头都填不满，这时切到别的后端或让多块模式兜底更划算。multi_query_tokens（投机解码）和 MLA 则有直接豁免：前者只有 XQA 支持，后者当前一律优先 XQA-MLA。

编译期特化：三个内核，一张宏表

XQA 的源码目录 cpp/kernels/xqa/ 里只有三份设备端实现：

三份文件共享同一套头文件（defines.h、mha.h、utils.cuh、barriers.cuh、tma.h 等），差别全在编译期宏。以 defines.h 为例，一个内核实例的形状由下面这些宏决定：

1
// allowed values are multiples of 16 in range [16, 256]（MLA 例外为 576）
2
#define HEAD_ELEMS 128          // 每头维度 D_h
3
#define HEAD_GRP_SIZE 8         // GQA 组大小 G
4
#define IS_MLA (HEAD_GRP_SIZE == 128 && HEAD_ELEMS == 576)
5
#define BEAM_WIDTH 1            // beam search 宽度，允许 1 或 4
6
#define SPEC_DEC 0              // 是否走投机解码验证路径
7
#define M_TILESIZE 32           // 投机路径下 HMMA 的 Q 行块（QGMMA 为 64）
8
#define CACHE_ELEM_ENUM 2       // 0: half/bf16  1: int8  2: fp8_e4m3
9
#define TOKENS_PER_PAGE 32      // 0 表示非分页（连续）KV
10
#define SHORT_SEQ_OPT 1         // 仅 mha.cu：短序列优化
11
#define SLIDING_WINDOW 0
12
#define ENABLE_PDL 2            // programmatic dependent launch

模板参数化的传统做法是「对每种组合都实例化一份」，结果是 cubin 数量爆炸——$D_h$ 有十余种取值、$G$ 有数种、KV 精度三种、页大小五种、beam 两种、投机两种。XQA 的选择是 运行时 JIT：源码以字符串形式打进 TensorRT-LLM 的共享库，第一次遇到某个形状组合时调用 NVRTC 现场编译，之后把 cubin 缓存下来复用。

compileEngine.cpp 负责把 XQAParams 翻译成 tllmXqaJitContext，再由 nvrtcWrapper.cpp 翻译成 -D 宏和编译选项。关键逻辑有三个：

1
bool const isLinearSpecDec = mXqaParams.multi_query_tokens && !mXqaParams.is_spec_dec_tree;
2
bool const useQGMMAKernel = supportConfigQGMMA(mXqaParams, mSM, true) && !isLinearSpecDec;
3
...
4
/*kernel_type=*/mXqaParams.isMLA() ? TLLM_XQA_JIT_MLA
5
                                   : (useQGMMAKernel ? TLLM_XQA_JIT_QGMMA : TLLM_XQA_JIT_HMMA),

第一处：线性投机解码（链式草稿 token，不是树状）不走 Hopper 的 QGMMA 内核。注释里写明了原因——QGMMA 路径的注意力在全原生 FP8 下计算，包括 $QK^\top$ 的输出；在链式投机验证这种「每个草稿位置都要被精确复核」的场景里，FP8 的 $QK^\top$ 误差会让本来合法的 token 被拒绝，等于白白损失了接受率。所以这里退回到 HMMA，并用更保守的数值路径。

第二处：Hopper 内核默认把 RoPE 也吞进注意力 kernel 里。appliesRoPEInXqaKernel() 要求 RoPE 维度是 16 的倍数、且属于 kLONG_ROPE、kROPE_GPT_NEOX 或 kROPE_GPTJ 之一，满足时才打开 USE_INPUT_KV=1 与 ROPE_STYLE，否则仍走独立的 invokeQKVPreprocessing 对 Q 做 RoPE。这样做省掉一次单独的 kernel 发射和一遍 Q 的读写，代价是内核里多一段旋转逻辑——这正是 decode 这类「launch 开销占比高」的场景最在意的。

第三处：nvrtcWrapper.cpp 推导 M_TILESIZE 的规则分三种：非投机时直接等于 GQA 组大小；投机且走 QGMMA 时固定为 64；投机且走 HMMA 时取 q_seq_len × G，不足 16 按 16、否则 32。把 TOKENS_PER_PAGE 直接映射成宏，SLIDING_WINDOW 在「树形投机」时强制为 0。编译选项固定为 -dw（忽略告警）、--use_fast_math、-default-device，Hopper 用 -arch=sm_90a（需要 wgmma 与 TMA 的 a 后缀特性），SM120 与 SM121 共用 -arch=sm_120f。

这套 JIT 的设计动机值得展开。decode 注意力对形状极其敏感：D_h=128 与 D_h=256 的 SMEM 布局、K 头的切分数、每瓦片的元素数全都不一样；G=8 与 G=32 直接决定 $M$ 方向的行数和 warp 分工；KV 是 FP8 还是 BF16 决定了是否需要转换与寄存器预算。如果预编译全部组合，一个 libtensorrt_llm.so 里要塞进几百份 cubin，加载时间和镜像体积都吃不消。反过来，纯运行时编译如果不缓存，每个新形状首次出现都要现场编译一遍（源码规模在数万行量级），对需要快速拉起新模型的服务来说是很明显的冷启动代价。XQA 的做法是 按 hash key 缓存：XQAKernelRuntimeHashKey 把 KV 数据类型、头维度、GQA 组大小、beam、M 瓦片、页大小、是否分页、是否投机、是否 FP8 输出、位置编码类型与 RoPE 维度打包成一个整数，CubinObjRegistry 按此查表，命中就直接发射。

Ampere 路径：一个 CTA 里的两类 warp

mha.cu 的 CTA 形状在编译期就固定下来：

1
// x: horizontal stacking for cta horizontal tile size
2
// y: vertical stacking for cta vertical tile size
3
// z: must be 2 for warp specialization.
4
constexpr uint3 ctaShapeInWarps = {4, 1, 2};

8 个 warp、256 个线程。z 方向上的两个平面承担两种角色：

• warpIdx.z == 0（4 个 warp）：负责搬数据并执行 GEMM0，即 $S = Q K^\top$。它们从全局内存（或分页池）把 K 瓦片搬进 SMEM，调用 mma.sync 得到 logits，再在 warp 内做行最大值、在线 softmax 的缩放。
• warpIdx.z == 1（4 个 warp）：执行 GEMM1，即 $O = P V$。它们把 softmax 后的概率矩阵 $P$ 与 V 瓦片相乘，并把结果写回输出。

这就是 FlashAttention 系的「两段 GEMM 切在两组 warp 上」在 decode 场景的落地。中间的 $S$ 矩阵从 GEMM0 到 GEMM1 的传递走 SMEM：GEMM0 warp 把结果写进 smem.x 的某个缓冲槽，用一个屏障通知 GEMM1；GEMM1 读完后，同一个缓冲要被 GEMM0 的下一瓦片复用，于是还要一个反向屏障。代码里把这一对循环屏障定义成：

FlashAttention-2 前向的 warp 工作划分可以作为参照：不同 warp 分别负责注意力矩阵的不同条带，靠共享内存交换中间结果。XQA 把这个原则换成了更粗的切法——不按矩阵条带切，而按 GEMM 阶段 切：4 个 warp 只做 $QK^\top$ 与 softmax，另外 4 个只做 $PV$ 与输出。来源：FlashAttention-2 论文（arXiv:2307.08691）Figure 3。

1
CtaBarrierPair xBarriers[ctaShapeInWarps.y][ctaShapeInWarps.x];

CtaBarrierPair 的语义是「一个表示生产者已就绪，一个表示消费者已清空」。用双向屏障而不是简单的 __syncthreads()，是因为两个角色执行完全异步：GEMM0 可能已经在做第 $n+1$ 块的 $QK^\top$，而 GEMM1 还在算第 $n$ 块的 $PV$。要让流水不互相等待，必须让「槽位 $n$ 可写」和「槽位 $n$ 可读」各自独立同步。

1
constexpr uint32_t preferedKHeadPartBytes = 64;   // SM86/89/120: 64 字节
2
constexpr uint32_t preferedKHeadPartBytes = 128;  // SM80/87/90: 128 字节
3
constexpr uint32_t cacheVTileSeqLen = 32;         // 同上两组

1
constexpr uint32_t nbPagesPerWarpTile
2
    = (warpTile.x <= tokensPerPage ? 1U : exactDiv(warpTile.x, tokensPerPage));

1
constexpr bool persistentQ = paddedInputHeadBytes * ctaTile.y <= (16u << 10);
2
static_assert(persistentQ);

1
#define GRP_LOAD_V (CACHE_ELEM_ENUM != 0) || (HEAD_ELEMS == 256 && USE_PAGED_KV_CACHE && BEAM_WIDTH > 1)

Hopper 路径：从两段 GEMM 到三角色流水

Hopper 上换了计算单元，一切都要重新组织。mha_sm90.cu 的 warp 数量、角色和缓冲结构都与 Ampere 不同：

1
constexpr uint32_t gmmaWarpsPerGrp = 4;
2
constexpr uint32_t gemm0NbGmmaGrps = 1;
3
constexpr uint32_t gemm0NbWarps = gmmaWarpsPerGrp * gemm0NbGmmaGrps;   // 4
4
constexpr uint32_t gemm1NbGmmaGrps = 1;
5
constexpr uint32_t gemm1NbWarps = gmmaWarpsPerGrp * gemm1NbGmmaGrps;   // 4
6
constexpr uint32_t nbIOWarps = 4;
7
constexpr uint32_t nbWarps = gemm0NbWarps + gemm1NbWarps + nbIOWarps; // 12

CTA 是 blockDim = (128, 1, 3)，共 12 个 warp。三类角色：

这套 Hopper 内核的适用范围比 Ampere 路径窄得多。supportConfigQGMMA() 额外要求：SM == kSM_90、输入是 FP16/BF16、KV 是 FP16/BF16/FP8，并且 非 skip-softmax 时 KV 必须是 FP8——也就是说，BF16 KV 的普通解码默认不会走 Hopper 内核，除非打开 SKIP_SOFTMAX_ATTN；它还不接受 beam search（beam_width != 1 直接拒绝）。设备端 mha_sm90.cu 的实现也被 #if __CUDA_ARCH__ >= 900 && BEAM_WIDTH == 1 包住，其他组合在编译期落到 trap。这看起来苛刻，实际是因为 Hopper 路径的设计目标就是「FP8 KV + 单序列 GQA 解码」：KV 压到 FP8 后，每字节带宽的产出翻倍，才值得为 TMA + wgmma 重写一整套数据流。BF16 KV 的收益不足以覆盖重写成本，所以留给 Ampere 内核。

与 Ampere 最大的不同在于 IO 从计算里拆出来了。Ampere 的 GEMM0 warp 亲自 cp.async 搬 K，搬的时候靠 ldgsts 的异步属性与 MMA 重叠；Hopper 的 TMA 本身就是一次异步事务（transaction），需要一个线程发起并提交 barrier 的预期字节数，而 wgmma 又是 warpgroup 级别的指令，要求同一个 warpgroup 的 128 个线程保持一致的执行节拍。把搬运交给独立的 IO warp，计算侧就能专注于「等 barrier、发 wgmma、做 softmax」，搬运侧可以同时处理多个缓冲区。

SharedMem 的布局也体现了双缓冲的规模：

1
static constexpr uint32_t nbKBuf = 2;
2
static constexpr uint32_t nbVBuf = 2;   // SKIP_SOFTMAX_ATTN 时为 3
3
static constexpr uint32_t nbXBuf
4
    = 2 * (gemm0CtaTileNbTokens >= gemm1CtaTileNbTokens ? 1 : exactDiv(gemm1CtaTileNbTokens, gemm0CtaTileNbTokens));

K 与 V 各双缓冲；X（softmax 后的 $P$）按「每 X 缓冲对应几个 V 缓冲」的比例扩展。一个很实用的技巧是 X、V 缓冲与输出暂存共用同一块 SMEM：

1
union ReusedXVOutSwizzleBuf
2
{
3
    struct XV { XBuffer x; VBuffer v; /* ... */ } xv;
4
    OutSwizzleBuf outSwizzle;
5
} reusedXVOutSwizzleBuf[nbXBuf];

$P$ 与 $V$ 只在 GEMM1 结束前有用，而输出重排（把累加器寄存器按目标布局写进 SMEM 再 TMA 出去）只发生在 GEMM1 之后。两者生命周期不重叠，所以可以叠在同一地址上，用 union 复用。文件里还有一个断言 static_assert(sizeof(OutSwizzleBuf) <= sizeof(XV))，保证复用不会溢出。这类「显存之外的显存」抠法在 Hopper 内核里很常见——SMEM 是 decode 注意力最紧的资源之一，每省一份缓冲就多一份流水深度。

1
#define SWAP_AB (!SPEC_DEC)   // 默认投机解码时不开，其余解码全部开启

1
bool const disableSkipForShortSeq = (cacheSeqLen < skipSoftmaxThresholdScaleFactor);
2
float const skipSoftmaxThreshold = disableSkipForShortSeq ? 0.0f : skipSoftmaxThresholdScaleFactor / cacheSeqLen;

1
auto initialRowMax = ThrdRegRowMax::filled(safeInitRowMax);
2
if (attentionSinks != nullptr)
3
{
4
    for (uint32_t i = 0; i < initialRowMax.size; i++)
5
    {
6
        uint32_t const rowOffset = warp_size * i + laneId();
7
        if (rowOffset < (SPEC_DEC ? warpTile.y : headGrpSize))
8
        {
9
            uint32_t const srcOffset = SPEC_DEC ? rowOffset % headGrpSize : rowOffset;
10
            initialRowMax[i] = fmaxf(initialRowMax[i], attentionSinks[headGrpSize * idxHeadGrp + srcOffset]);
11
        }
12
    }
13
}

分页 KV：页号怎么变成 TMA 坐标

XQA 对分页 KV 最底下一层的抽象是 KVCacheList：它把「全局页池 + 每条序列的页表」包在一起，向内核提供两样东西——给 TMA 用的 tensor map，以及给普通 gather 用的 getHead()。两条路径并存，是因为 Ampere 没有 TMA、必须手工拼地址，而 Hopper 可以把寻址交给 TMA 硬件。

1
// VLLM Layout
2
uint64_t const globalDims[] = {headElems, nbKHeads, tokensPerPage, 1U << 31};
3
uint32_t const headBytes = elemBytes * headElems;
4
uint64_t const globalStrides[] = {headBytes, headBytes * nbKHeads, headBytes * nbKHeads * tokensPerPage};
5
uint32_t partBytes = partElems * elemBytes;
6
uint32_t const boxDims[] = {partElems, 1, mha::min(tokensPerPage, nbTokensPerTile), 1};

1
tma::loadAsync(&dst(tokensPerPage * i, 0), tensorMap,
2
    DimsLE<4>{partElems * idxPart, idxHeadGrp, 0, (uint32_t) pages[i]}, bar);

1
uint32_t const idxPageBeg = gemm0CtaTileNbTokens >= tokensPerPage
2
    ? nbPagesPerTile * idxTile
3
    : idxTile / exactDiv(tokensPerPage, gemm0CtaTileNbTokens);
4
for (uint32_t i = 0; i < nbPagesPerTile; i++)
5
{
6
    uint32_t const idxPage = idxPageBeg + i;
7
    auto const page = (idxPage >= nbSkipLeadingPages && idxPage < nbPages)
8
        ? cacheList.kvCachePageList[baseOffset + idxPage]
9
        : kBAD_PAGE_INDEX;
10
    if (warpElectSync()) pages[i] = page;
11
}

1
constexpr uint32_t nbPagesPerWarpTile = (warpTile.x <= tokensPerPage ? 1U : exactDiv(warpTile.x, tokensPerPage));
2
using KCachePageIndices = Vec<KVCachePageIndex, nbPagesPerWarpTile>;

1
// XQA 原生布局（K/V 共池）
2
HeadPtr<GMemCacheHead const, tokensPerPage, nbPagesPerWarpTile> const src{
3
    cacheList.pool, pageIdx, nbKHeads, idxHeadBeg};

FP8 / INT8 的 KV：转换到底放在哪儿

XQA 支持 FP16、BF16、FP8(E4M3)、INT8 四种 KV 缓存类型，计算精度则取决于路径：Ampere 先转成 FP16/BF16 再做 mma，Hopper 的 FP8 路径直接在 FP8 上做 wgmma。差异体现在 CACHE_ELEM_ENUM 宏：

1
float const warpRank = warpIdx.x;
2
constexpr float xScale = 1.f / kE4M3_MAX;
3
#if LOW_PREC_OUTPUT
4
float const oScale = rcpOutScale[0];
5
#else
6
constexpr float oScale = 1.F;
7
#endif
8
float const xvoScale = xScale * (isKVCacheQuantized ? kvCacheScale[0] : 1.f) * oScale;

1
float const qkScale = qScale * (isKVCacheQuantized ? kvCacheScale[0] : 1.f) * rsqrtf(validElemsPerHead);

1
auto const kSlice = [&]()
2
{
3
    if constexpr (mha::is_same_v<InputElem, KElemType>) return kSliceOrig;
4
    else if constexpr (mha::is_same_v<KElemType, int8_t> || mha::is_same_v<KElemType, __nv_fp8_e4m3>)
5
    {
6
        Array2D<...> ret;
7
        for (m, n, i, j)
8
        {
9
            auto const data = convertKCacheWordToF16<InputElem, KElemType>(kSliceOrig(m, n).data[i][j]);
10
            ret(m, n).data[i][j * cvtExp] = data[0];
11
            ret(m, n).data[i][j * cvtExp + 1] = data[1];
12
        }
13
    }
14
}();

1
__device__ inline Vec<uint32_t, 2> convertKCacheWordToF16(uint32_t i8data)
2
{
3
#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 890)
4
    if constexpr (mha::is_same_v<InputElem, half> && mha::is_same_v<CacheElem, __nv_fp8_e4m3>)
5
    {
6
        uint16_t(&src)[2] = reinterpret_cast<uint16_t(&)[2]>(i8data);
7
        uint32_t(&dst)[2] = reinterpret_cast<uint32_t(&)[2]>(ret);
8
        asm("{\n"
9
            "cvt.rn.f16x2.e4m3x2 %0, %2;\n"
10
            "cvt.rn.f16x2.e4m3x2 %1, %3;\n"
11
            "}\n" : "=r"(dst[0]), "=r"(dst[1]) : "h"(src[0]), "h"(src[1]));
12
        return ret;
13
    }
14
#endif
15
    // 其他组合：逐元素标量转换
16
    for (uint32_t i = 0; i < 4; i++) dst[i] = InputElem(src[i]);
17
    return ret;
18
}

多块模式：把一条序列沿 KV 切成几段

前面反复提到「并行度不足」。TP8 部署的 70B GQA 模型每卡只有 1 个 KV 头，batch 32 时也只有 32 个 CTA，在 132 个 SM 的 H100 上不到 1/4 个 wave。更极端的是单条长上下文请求：每卡只有 1 个 CTA，132 个 SM 里 131 个在看戏——多块模式就是为这种「CTA 数量远小于 SM 数量」的场景准备的。

XQA 的应对方式是 multi-block：沿着 KV 的序列维把一条序列切成 $n$ 个子序列，每个子序列由一个 CTA 负责，各自独立算出「局部 softmax 结果」，最后由一个 CTA 做合并。这本质上是 FlashDecoding 的思路——解码的 $M$ 方向太小，就把 $N$ 方向拆开。可参考 FlashAttention 完全拆解（二）：序列维并行、Warp 工作划分与 FlashDecoding 对 split-KV 归约的讲解。

FlashAttention-2 的并行划分示意：前向时每个 worker 负责注意力矩阵的一段行（一段 query），反向后每个 worker 负责一段列。XQA 的多块模式采用同一类思路，但方向相反——decode 只有一行 query，于是沿 KV 的列方向切，把一条序列分给多个 CTA，最后做 split-KV 归约。来源：FlashAttention-2 论文（arXiv:2307.08691）Figure 2。

1
float const factor = 0.25f;
2
return mha::min<uint32_t>(
3
    mha::max<uint32_t>(1U, (uint32_t) round(prop.multiProcessorCount * 3 / (batchSize * nbKHeads) * factor)),
4
    divUp(maxSeqLen, gemm0CtaTileNbTokens));
5

6
// mha.cu
7
return std::min<uint32_t>(
8
    std::max<uint32_t>(1U, prop.multiProcessorCount / (batchSize * nbKHeads)), divUp(maxSeqLen, ctaTile.x));

1
multiBlockCount = historyLength / kMinHistoryTokensPerBlock;            // 128
2
multiBlockCount = std::min(multiBlockCount,
3
    (int) std::round(std::sqrt(multiBlockCount * 8.F)));               // 防止归约项主导
4
multiBlockCount = std::max(multiBlockCount, 1);
5
double waveCount = (double) batchSize * numKVHeads * multiBlockCount / (double) multiprocessorCount;
6
double adjFactor = waveCount / (double) kTargetWaveFactor;            // 8
7
if (adjFactor > 1.0) multiBlockCount = floor(multiBlockCount / adjFactor);
8
multiBlockCount = std::max(multiBlockCount, 1);
9
multiBlockCount = std::max(std::min(multiBlockCount, maxNbSubSeq / batchSize), 1);

1
mRowSumMax = segmenter.newSeg<ColWiseVec>(nbChunks, alignment);   // ColWiseVec = Vec<SumMax, ctaNbValidQHeads>
2
mTokens    = segmenter.newSeg<Vec<IOHead, ctaNbValidQHeads>>(nbChunks, alignment);

1
if (sumMax.max > state.max)
2
{
3
    float const scale = expf(state.max - sumMax.max);
4
    state.max = sumMax.max;
5
    state.sum = state.sum * scale + sumMax.sum;
6
    state.acc = state.acc * scale + data * sumMax.sum;
7
}
8
else
9
{
10
    float const scale = expf(sumMax.max - state.max);
11
    state.sum = state.sum + sumMax.sum * scale;
12
    state.acc = state.acc + data * (sumMax.sum * scale);
13
}

1
uint32_t const lastOld = nbSubSeq - 1;
2
uint32_t const idxSemaphore = idxSeq * nbInputSeqSplit + idxInputSubSeq;
3
auto const pSemaphore = &semaphores[idxSemaphore];
4
asm volatile("atom.acq_rel.gpu.global.inc.u32 %0, [%1], %2;\n" : "=r"(old) : "l"(pSemaphore), "r"(lastOld));
5
smemIsLastCta = (old == lastOld);

投机解码：XQA 的第二个工作负载

启用投机解码后，decode step 不再只算一个 query token，而是一次性验证一段草稿 token（树形或链式）。defines.h 把 SPEC_DEC=1 打开时，内核形态变化很大：

1
#if SPEC_DEC
2
using MaskType = uint32_t;
3
#ifndef M_TILESIZE
4
#define M_TILESIZE 32
5
#endif
6
#endif

关键变化有四处：

1. inputSeqLen 不再是 1：mha.h 里的 constexpr uint32_t inputSeqLen = 1; 只在非投机时成立，投机时等于 generation_input_length，即每条请求的草稿长度。
2. 网格的 y 维变成「头 × 草稿 token 块」：nbTokenBlocksPerGrp = divUp(qSeqLen * headGrpSize, m_tilesize)，把「草稿 token × GQA 头」压成一维，再由 warpTile.y = roundUp(rowsPerBlock, 16) 分块。这样 $M$ 方向可以填到 16 或 32 行，而不是 8 行。
3. 需要树形掩码：IS_SPEC_DEC_TREE=1（默认）时，草稿 token 之间不是链式关系，而是一棵树，每个 token 只能看到它在树上的祖先。内核在 warpGrpApplyMask 里逐行读 mask（mask 是一个 uint32_t*，每个草稿 token 一行 bitmask），把非祖先位置的 logit 置为 safeInitRowMax（一个大负数），从而被 softmax 忽略。
4. SWAP_AB 默认关闭：mha_sm90.cu 里默认定义 #define SWAP_AB (!SPEC_DEC)，树形验证时回到 K 在 B 位、Q 在 A 位的标准布局（固定草稿长度的 SPEC_Q_SEQ_LEN 变体是例外，它强制开启 SWAP_AB）。树形掩码要按 query token 逐行处理，而 SWAP_AB 下 $S$ 的每一列才对应一个 query token，二者方向不一致。

CompileEngine 里的第一处判断正是关于这里的：

1
bool const isLinearSpecDec = mXqaParams.multi_query_tokens && !mXqaParams.is_spec_dec_tree;
2
bool const useQGMMAKernel = supportConfigQGMMA(mXqaParams, mSM, true) && !isLinearSpecDec;

链式（linear）草稿在 Hopper 上不走 QGMMA，而树形草稿走。前者的原因已经说过：全原生 FP8 的 $QK^\top$ 会损失验证精度；后者的原因要从数据流上理解：树形验证的 $M$ 方向被多个草稿 token 填满，wgmma 的 $M=64$ 粒度不再浪费，FP8 带来的带宽节省变得更值；而且 QGMMA 的 K 瓦片是 64 个 token 一块，树形草稿共享的 KV 前缀能在一次 TMA 里被多个分支复用。

mha.cu 的网格选择也按投机与非投机分叉：

1
#if SPEC_DEC
2
    const uint32_t nbTokenBlocksPerGrp = divUp(qSeqLen * headGrpSize, rowsPerBlock);
3
    dim3 const dimGrid{nbSubSeqPerSeq, nbKHeads * nbTokenBlocksPerGrp, batchSize};
4
#else
5
    dim3 const dimGrid{nbSubSeqPerSeq, nbKHeads, batchSize};
6
#endif

非投机的 dimGrid.y 是 nbKHeads；投机的 dimGrid.y 是 nbKHeads × nbTokenBlocksPerGrp，让每个草稿 token 块与每类头各自独立，互不依赖。投机路径同样可以开多块，但启发式换了：computeMultiBlockCountSpecDec() 只在 batch 不超过 32、历史长度不小于 2048 时才调优；它先按「单块 CTA 数不足一个 wave」的条件算出可填充的块数，取不超过一波 SM 的 2 的幂并限制在 64 以内，再按历史长度右移收缩，最后断言「块数 × 单块网格」不超过 SM 数——避免多开一个 wave 导致尾延迟。

和其他路径怎么配合：为什么 XQA 不是唯一答案

把 XQA 放回 TensorRT-LLM 的注意力后端体系，就能看到一条完整的分工链：

1. XqaDispatcher 构造函数先看是不是 Blackwell（isSM100Family()），是的话直接创建 TllmGenFmhaRunner，把整个注意力路径交给 TRTLLM-Gen FMHA；只有非 SM100 才构造 DecoderXQARunner。
2. 不是 Blackwell，就用 DecoderXQARunner 判断「这个形状该不该走 XQA」，不行就退回 masked MHA。
3. 是 MLA 模型时，computeMultiBlockCountForMLA() 直接返回 1；注释只写了「disable multi-block for MLA kernel for now」，没有给出原因。MLA 的 KV 布局、归约与页表都另写了一套，当前版本尚未把多块接进去。
4. FlashInfer 后端完全走自己的路径，与 XQA 无关——它有自己的分页 KV 表示与 JIT 模板，可参考 FlashInfer 完全拆解：把分页 KV、动态批处理与定制注意力编译成服务级 Kernel。

这个分工背后有一个很强的实用主义立场：解码注意力没有一劳永逸的最佳内核。XQA 在「GQA + FP8 KV + 中长上下文」的 Hopper 解码上是首选；但 BF16 KV 的 Hopper 解码可能仍走 HMMA；在 Ampere 上是否优于 MMHA 取决于 batch 与上下文能否填满 SM；在 SM120 且 FP8 条件满足时，MLA 会优先走 XQA-MLA；在 SM100 上它彻底退场。理解这条边界，比记住「XQA 很快」更重要。

一次 decode step 的完整时序

把前面几节串起来，XQA 在一次 decode 前向里的执行顺序是这样的：

1. 准入与编译。XqaDispatcher::prepare() 在第一次遇到某种形状时调用 DecoderXQARunner::prepare()：shouldUse() 先过 supportConfig() 与 mayHavePerfGain() 两道关，再按 CubinObjKey 查 CubinObjRegistry；未命中则以 CompileEngine 调 NVRTC 现场生成 cubin 并注册。
2. QKV 预处理。run() 里若内核不接管 RoPE（Ampere 一律、Hopper 仅在 RoPE 维度不满足对齐时），先调 invokeQKVPreprocessing：应用 QKV bias 与 RoPE，把处理后的 Q 写进 scratch 的 xqa_q_input_ptr，并把 $1/\sqrt{d_k}$ 这类缩放比例作为标量参数交给注意力（preprocessingParms.fmha_host_bmm1_scale = 1.0f / (sqrtf(head_size) * q_scaling)）。XQA 的内核读的是 已经追加过新 token 的 KV 缓存，这一步的前提是上游已经把本步的 K/V 写进分页池。
3. 构造网格与参数。run() 按 isMLA / isSpecDec && isHMMAKernel / 常规三条分支组装 kernel 参数数组：TensorMap（Hopper）、页表指针、attentionSinks、kvCacheScale、semaphores、scratch、beamSearchParams。网格尺寸按第 2 节和第 10 节里讨论的规则决定。
4. 内核执行。Ampere：8 warp、两平面分工，cp.async 搬 K/V，在线 softmax。Hopper：12 warp、三角色分工，IO warp 用 TMA 搬 K/V，GEMM0/1 用 wgmma。多块模式下每个 CTA 只算一段 KV，写完 scratch 后对信号量做一次原子自增。
5. 归约与输出。单块模式直接写输出；多块模式由最后一个 CTA 读全部 scratch 做在线归约，加回 attention sink，最后除以分母写回。若 LOW_PREC_OUTPUT，输出是 FP8，需要用 rcpOutScale 还原。

这个时序里有两个值得记的点：

• preprocessing 与注意力是两个 kernel。在 Ampere 上，Q 的 bias、RoPE 与缩放准备都发生在 preprocessing 里，注意力只拿到「已经转好」的 Q 与一个标量缩放因子；Hopper 在满足条件时把 RoPE 合进注意力，preprocessing 这一步就省掉了。这也意味着在 Ampere 路径上，XQA 并不省掉那次预处理的开销，它省下的是注意力本身在 SM 上的时间。
• KV 的追加不在 XQA 里。新 token 的 K/V 由上游（preprocessing 或框架的 KV 写入路径）落进分页池，XQA 只从页表开始读。这条边界让 XQA 与任何 KV 管理策略解耦：无论页大小、是否前缀共享、是否 P/D 分离中从别的实例收到 KV，只要页表被正确填好，XQA 就能直接读。

工程取舍与局限

读完源码，几个值得记录的取舍：

JIT 换体积和延迟。按形状编译的好处是每个 cubin 都只服务它自己的 $D_h$、$G$、页大小和精度，寄存器分配、SMEM 布局全是最好情况。代价是首次形状出现时的编译延迟与 cubin 缓存的管理成本。TensorRT-LLM 用 CubinObjRegistry 做进程级缓存，并在编译时按 hash 复用，但多进程部署时每个进程都会各自 JIT。仓库里还保留了 gen_cubins.py，用于把测试常用组合预先生成 cubin 再嵌入库，算是一种「预编译 + JIT」的折中。

确定性与性能的二选一。CTA_ROW_MAX_BACKWARD_METHOD 选 1 而非 4，损失约 8%；但换来的是逐位可复现，这在 SLO 严格、需要结果可比的服务里是硬需求。

多块的归约项。多块的收益在长上下文、小 batch 时最大；但块数过多会让归约项（读 scratch、合并 softmax）变成新的瓶颈。sqrt(8k) 的砍法是经验值，论文或文档里没有给出推导，实际部署时按硬件与负载调整 XQA_NB_SUB_SEQ 环境变量更可靠。

FP8 精度与接受率。QGMMA 路径上链式投机被禁用，说明全 FP8 的注意力在验证环节的误差不可忽视。如果模型输入 FP16/BF16、KV FP8，同时启用链式草稿，就会落到 HMMA：FP8 的带宽收益仍在（$QK^\top$ 在 FP16 中完成），但 Hopper 的 TMA 与 wgmma 优势用不上，性能曲线与纯 FP8 的 QGMMA 路径并不相同。

sink 与滑窗的组合。SLIDING_WINDOW 与 sink 都支持，但 supportConfigCommon() 要求 sink_token_length == 0，这意味着 XQA 目前不支持带 sink token 的流式生成——内核里有 attentionSinks 参数，但那是另一个含义（attention sink 权重，不是注入的 sink token）。这两者容易混淆，读源码时必须区分。

小结

XQA 是 TensorRT-LLM 在 Ampere/Hopper 上把 decode 注意力「做到极致」的一条路径，也是 SM120 上 MLA 解码的当前实现。它的核心并不在某个非常巧妙的数学技巧，而在三件事的工程组织：

1. 形状特化到编译期：用 NVRTC 按 $D_h$、$G$、页大小、KV 精度、投机模式现场编译 cubin，每个实例都是最优寄存器与 SMEM 布局。
2. 角色分工：Ampere 把 GEMM0 与 GEMM1 分到两个 warp 平面，靠双向屏障做流水；Hopper 把搬运拆成独立 IO warp，用 TMA + wgmma 做三角色异步。
3. 补足并行度：多块模式沿 KV 切序列，用 scratch 存局部和-最大值，信号量选最后一个 CTA 归约；投机路径则把草稿树压进 $M$ 方向，让 wgmma 的粒度对齐。

最后回到开头的问题：XQA 凭什么比 MMHA 快？对 GQA 解码，它把「每读一份 KV 喂 $G$ 个 Q 头」这个复用关系做到了数据布局层、SMEM 层与指令层；凭什么在某些形状下不如 MMHA？当每卡的「batch × 本地 KV 头数」太小、KV 太短、或者模型是 MHA 时，它的特化布局没有复用可言，额外的 JIT 与参数组装反而成了负担——这正是 mayHavePerfGain() 与 supportConfigCommon() 两道关卡存在的理由。

参考资料

1. TensorRT-LLM XQA kernel source (cpp/kernels/xqa)
2. TensorRT-LLM Attention docs: XQA optimization support matrix
3. decoderXQARunner.cpp: shouldUse / grid assembly
4. decoderXQAImplJIT/compileEngine.cpp: JIT macro translation
5. TensorRT-LLM kernels: xqaDispatcher.cpp
6. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
7. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
8. GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints
9. CUDA Tensor Memory Accelerator (TMA) programming documentation