# gfx936 DCU上实现INT8 QK MMAC：分页访存、Fragment映射与GQA适配 ## 前言 本文是系列第二篇。第一篇[《gfx936 DCU上实现INT8 KV与INT8 MMAC Attention推理优化》](https://blog.csdn.net/Foreverahead/article/details/162978554?spm=1001.2014.3001.5501)介绍了完整数据流，本文聚焦 Attention 的第一次矩阵乘法 `QK^T`。 把 K Cache 存成 INT8 并不代表 QK 已经使用低精度计算。最容易实现的一种路径是从 Cache 读取 INT8 K，在寄存器中转成 BF16，再调用原来的 BF16 MMAC。我们在真实 Qwen tile 上测过，这条路线只有 BF16 QK 的 `0.71x-1.00x`，读取量虽然下降，转换和 Scale 开销却把收益吃掉了。 最终采用的路线是：Query 也做动态 INT8 量化，让 `Q_int8×K_int8^T` 直接进入 gfx936 的 `v_mmac_i32_16x16x32_i8`，在 INT32 中累加，再恢复 Query Scale、Key Scale 和 softmax Scale。实现过程中真正困难的不是调用一条 MMAC 指令，而是让分页 K Cache、Query fragment、Scale 下标和 softmax token 使用同一套 lane 映射。 本文实验环境：单张 gfx936 DCU、DTK 26.04、Qwen3.5-27B。对应代码主要位于 `csrc/rocm/attention.cu`：Decode 使用 `paged_attention_int8_kv`，带历史的分块 Prefill 使用 `gfx936_int8_prefix_attention_kernel`，Python 入口位于 `vllm/v1/attention/backends/rocm_aiter_fa.py`。 本文代码仓库：[https://github.com/sinpeyw/qwen3.5-vllm-dcu-optimization](https://github.com/sinpeyw/qwen3.5-vllm-dcu-optimization) ## 1. INT8 QK的量化和Scale恢复 Attention Score 为： ```text S = Q × K^T / sqrt(d) ``` 对 Query 和 Key 分别做对称动态量化： ```text Q ≈ Qq × sQ K ≈ Kq × sK ``` 代回原式可得： ```text S ≈ MMAC_INT8(Qq, Kq^T) × sQ × sK / sqrt(d) ``` `Qq`、`Kq` 是 INT8，MMAC 输出是 INT32 accumulator。转成 FP32 后，每个 score 需要乘： ```text post_scale = sQ(query_token, query_head) × sK(key_token, kv_head) × softmax_scale ``` 这里有一个容易忽视的区别：Query Scale 随当前 Query token 和 Query head 变化，Key Scale 则随历史 key token 和 KV head 变化。最终 K Scale 采用 per-token、per-KV-head 粒度，因此一个 QK tile 中不同列对应的 Scale 也不同，不能在整块矩阵乘法结束后只乘一个统一系数。 Scale 还必须在 score 进入 softmax 之前恢复。若先对 INT32 accumulator 做 softmax，再补乘 Key Scale，改变的就不只是近似精度，而是不同 key token 之间的相对概率。最终实现不生成全局 Score tensor，而是在寄存器中恢复每个 score，随后直接进入 online softmax。 ## 2. Q和K如何映射到gfx936 MMAC ### 2.1 Q放在operand A，K放在operand B 本文使用的整数矩阵指令为： ```text v_mmac_i32_16x16x32_i8 ``` 一个 wave64 协同完成： ```text A: 16 × 32, INT8 B: 32 × 16, INT8 C: 16 × 16, INT32 ``` 在这个布局中，operand A 对应输出矩阵的行，operand B 对应输出矩阵的列。QK 的自然映射是： ```text A = Q，输出行对应 Query row B = K，输出列对应历史 token ``` 我们早期也考虑过执行 `MMAC(K,Q)` 再转置。数学上两者等价，但这种顺序会让输出归属、Key Scale 读取和 softmax 写入都变复杂。最终 CPU fragment 模型和 GPU 内核都固定为 Q 在 A、K 在 B。 对于 Qwen3.5 的 `head_dim=256`，INT8 指令每次推进 K 维 32，因此只需： ```text 256 / 32 = 8 个 K-step ``` 相同输出块的 BF16 指令每次推进 K 维 16，需要 16 个 K-step。这个差异只属于矩阵乘加部分，分页寻址、Scale 恢复和 softmax 不会同时减半。 INT32 累加也不存在溢出压力。量化范围为 `[-127,127]` 时，完整 256 维点积的绝对上界为： ```text 256 × 127 × 127 = 4,129,024 ``` 它距离 `2^31-1` 仍有约 500 倍余量。实际精度风险主要来自量化后 score 的相对扰动，而不是 INT32 饱和。 ### 2.2 MMAC输出不是四个连续token 一个 wave 有 64 个 lane，每个 lane 为 A、B 分别提供 8 个 INT8 元素，并持有 4 个 INT32 输出。下文把每个 lane 持有的输入或输出片段称为 fragment。 MMAC 的四个输出并不是四个连续列。令： ```text lane_col = lane_id % 16 ``` 一个 lane 实际持有的输出列为： ```text lane_col + {0, 4, 8, 12} ``` 例如 `lane_col=2` 时，对应的是第 `2、6、10、14` 列，而不是第 `2、3、4、5` 列。若共享内存写入、Key Scale 读取或尾部 mask 仍按连续四个 token 处理，内核可以正常结束，结果却会悄悄错位。 QK 主循环可以简化为： ```cpp int32x4 acc = {0, 0, 0, 0}; #pragma unroll for (int k_base = 0; k_base < 256; k_base += 32) { int8x8 q_frag = load_q_fragment(q_int8, k_base, lane); int8x8 k_frag = load_k_fragment(k_cache, page, slot, k_base, lane); acc = mmac_i32_16x16x32_i8(q_frag, k_frag, acc); } int lane_col = lane_id % 16; for (int i = 0; i < 4; ++i) { int col = lane_col + 4 * i; int key_token = token_base + col; score[i] = float(acc[i]) * q_scale * k_scale[key_token][kv_head] * softmax_scale; } ``` 真正的难点集中在 `load_q_fragment`、`load_k_fragment` 和 `key_token` 三者是否对应同一个矩阵位置。最终代码保留 MMAC 的寄存器顺序，到写入 softmax 工作区时再完成 4×4 转置。 ### 2.3 保守指令封装和流水指令封装 最终代码保留了两种 wrapper。用于独立验证和 Decode QK 循环的保守版本，在 MMAC 前后保留等待： ```cpp asm volatile( "s_nop 4\n\t" "v_mmac_i32_16x16x32_i8 %0, %1, %2, %0\n\t" "s_nop 4\n\t" : "+v"(acc) : "v"(a), "v"(b)); ``` 分块 Prefill 使用的流水版本只发出 MMAC，让展开循环中的 fragment load、Scale 计算和其他独立工作覆盖指令间隔。删除 `s_nop` 并不是普遍成立的优化；若调用点没有足够的独立指令，硬件相关等待仍可能重新出现，甚至产生时序或正确性问题。 因此我们同时检查 `dccobjdump` 中的真实指令、VGPR 与 scratch 使用量，以及完整 kernel 的实测时延。只看 C++ 源码里少了两行 `s_nop`，不能证明内核已经变快。 ## 3. 分页K Cache如何持续为MMAC供数 PagedAttention 中的 K 并不是一个连续的 `[token,head,dim]` 矩阵。逻辑 token 需要先通过 block table 找到物理 block，再通过 slot 找到页内位置。如果完全按照逻辑矩阵逐元素读取，MMAC 的 16 个输出列会产生大量分散访问。 第一篇已经介绍过 K Cache 的总体布局，这里只看 QK 如何消费它： ```text K cache: [block, kv_head, head_dim / 16, block_size, 16] ``` 加载一个 K fragment 时，内核依次完成： 1. 根据逻辑 key token 查询 block table； 2. 计算该 token 在物理页中的 slot； 3. 根据 `lane_k_base/16` 选择 head-dimension 小块； 4. 用 8-byte 或 16-byte 向量 load 取得 INT8 fragment； 5. 将 fragment 直接送入 MMAC。 这里不生成完整 BF16 K tensor。K 从分页 INT8 Cache 进入寄存器或 LDS 后，仍保持 INT8，直到参与矩阵乘法。这样才能同时保留两类收益：读取 K 的有效字节数下降，QK 又能使用原生 INT8 MMAC。 我们专门测过另一条看似保守的路线：读取 INT8 K，在寄存器里乘 Scale 并转成 BF16，再调用成熟 BF16 MMAC。它不需要全局 BF16 中间张量，但真实 Qwen tile 只有 BF16 对照的 `0.71x-1.00x`。原因是短 fragment 上的类型转换、Scale 和重排指令已经足以抵消读取量下降。 ## 4. Query量化只做一次 K 已经保存在 INT8 Cache 中，Q 则是当前步骤新生成的 BF16 数据。每个 Query row 需要计算： ```text sQ = max(abs(Q)) / 127 Qq = round(Q / sQ) ``` Decode Attention 会把长上下文分成多个 segment。若每个 segment 都重新量化 Q，同一条 256-element Query 会被反复读取、归约和舍入。分段越多，这项重复成本越明显。 最终实现把 Q 量化并入 RoPE/KV producer。一个 Query 只生成一份 INT8 数据和 Scale，所有 split-K segment 共同复用。带历史的分块 Prefill也采用相同原则：Query 量化发生在 Attention 之前，不放进每个历史分段内部。 为了判断这项成本的量级，我们曾单独测量处理全部 24 个 Query 头的成熟量化内核： | 当前 Query token 数 | Query量化时间 | | ------------------: | ------------: | | 4096 | `0.450 ms` | | 8192 | `0.896 ms` | | 16384 | `1.787 ms` | 这些是 Prefill/Extend 大块 Query 的测试，不是单 token Decode 的量化时间。和 20 ms 以上的长上下文 Attention 相比，它不是最大项，但也绝不是零成本。把它融合进 producer 的价值在于减少全局 Q 读写、kernel launch 和分段重复归约。 ## 5. Decode和Prefill不能共用一种行映射 Qwen3.5 的全 Attention 使用 24 个 Query 头和 4 个 KV 头，即 6:1 GQA。每个 KV 头由 6 个 Query 头共享，而 gfx936 MMAC 的输出块固定有 16 行。 ### 5.1 Decode：16行中只有6种独立Query Decode 每次只有一个新 token，无法用 token 维填满 16 行。最终 `wxp` 内核按照 Query 头编号对 6 取模，将 6 个 Query 头重复填入 16 个 MMAC 行。 16 行都会执行，但只有 6 行包含彼此不同的 Query，因此独立有效行比例为： ```text 6 / 16 = 37.5% ``` 这不等于只有 37.5% 的线程工作，而是 62.5% 的矩阵行重复了已有 Query。INT8 的 K-step 虽然减半，这部分重复计算仍然存在。比赛结束前，另一种 GQA6 紧凑打包布局没有完成完整正确性和端到端验证，因此最终提交接受了这一上限。 ### 5.2 Prefill：用Query token填满16行 带历史的分块 Prefill 有大量当前 Query token，可以采用更自然的映射：每个 Query 头单独处理，MMAC 的行对应当前输入块中的 Query token，KV 头通过 `query_head/6` 映射。 早期内核由一个 workgroup 处理 64 个 Query row，改进后由 8 个 wave 协同处理 128 个 Query row，使同一份 K/V tile 被更多 Query 复用。这里的 128 是 Query token 方向上的行数，不是 128 个 Query 头。 因此，Decode 与 Prefill 的共同点是都执行 INT8 QK，区别在于 M 维从哪里取得并行度： ```text Decode： M维来自GQA组内Query头，只有6种独立行 Prefill： M维来自当前Query token，可以填满MMAC行 ``` 这也是为什么一个在 Prefill 上利用率很高的 QK tile，不能直接套到单请求 Decode。 ## 6. QK算子实测结果 为了避免把其他算子的收益算到 QK 上，我们只保留两类 QK 对照： | QK路线 | 对照 | 实测结果 | 结论 | | ----------------------------- | ----------- | -----------------------: | -------------------------- | | INT8 K转BF16，再执行BF16 MMAC | 原生BF16 QK | `0.71x-1.00x` | 转换和Scale抵消读取收益 | | Query和Key直接执行INT8 MMAC | 原生BF16 QK | 大tile为 `1.55x-1.63x` | 长上下文QK有明确收益 | | 原生INT8 MMAC小形状 | 原生BF16 QK | 256-token附近接近持平 | kernel启动和固定开销占主导 | 这组数据说明，INT8 K Cache 必须和 INT8 QK 计算结合，单纯把 K 压缩后再恢复 BF16 很难拿到同样收益。同时，`1.55x-1.63x` 只覆盖 QK 内层矩阵块，不包含 Query 量化、online softmax、PV、KV producer 和输出写回。 完整 Attention 和端到端收益已经在第一篇给出。最终 QK 与 INT8 PV、融合 producer 和选择性 Prefill 一起进入正式提交. ## 7. QK正确性和适用边界 QK 位于 softmax 之前，一个不大的 score 误差也可能改变最终概率排序。我们采用了四层验证： 1. **CPU fragment模型**：重建 wave64 的 `16×16×32` 输入和输出，确认每个 lane 的四个 accumulator 对应哪些列； 2. **完整head_dim=256**：INT8 对照覆盖 8 个 K-step，BF16 对照覆盖 16 个 K-step，避免只计算一半维度； 3. **分页和Scale**：覆盖跨物理 block、非零 slot、per-token Key Scale、尾部 mask 和 GQA head mapping； 4. **运行时路径**：检查目标原生扩展、路由命中与真实 kernel 调用，而不是只看二进制中存在 MMAC 指令。 其中第二项来自一次真实教训：早期 BF16 控制组只加载了 128 维 Query，导致我们错误地把控制组问题归因成 INT8 误差。第三项也不能只验证连续 token，因为 MMAC 输出列本身是交织的。 最终这条专用路径明确检查 `head_size=256`、GQA ratio为6、无 sliding window和无 shuffle cache；Prefill 路径还要求 `Q24/KV4`，并排除 ALiBi 与 logits soft cap。页面布局、stride或模型结构不满足时直接拒绝启用，而不是带着错误映射静默运行。 INT8 QK更适合： - head dimension较大，例如256； - 历史上下文足够长，矩阵乘法不再主要受启动开销限制； - K已经按MMAC读取方向排布； - Query量化可以融合并跨segment复用； - MMAC行能由token或多个GQA group有效填充； - INT32结果可以在寄存器中恢复Scale并直接进入softmax。 不适合直接套用的情况包括： - 很小的QK tile； - 先完整反量化，再调用另一个BF16 kernel； - 每个split-K segment重复量化Query； - 只比较MMAC峰值，不检查GQA行利用率； - 没有证明服务实际执行目标内核。 ## 8. 结语 在 gfx936 上实现 INT8 QK，最关键的并不是把 Query 和 Key 转成 INT8，而是让分页 K Cache、Query fragment、Scale 下标和 softmax token严格对齐到同一套 MMAC 寄存器映射。 最终可以确认两点：原生 INT8 QK MMAC 在长 tile 上达到 `1.55x-1.63x`，而“INT8 K转BF16后继续使用BF16 MMAC”没有稳定收益。低精度数据只有一直保留到矩阵指令入口，才能同时兑现带宽和计算两方面的价值。 下一篇[《gfx936 DCU上实现INT8 PV MMAC：V Scale融合、概率量化与Fragment分组》](https://blog.csdn.net/Foreverahead/article/details/163176013?spm=1001.2014.3001.5501)将继续介绍更复杂的 INT8 PV。