gfx936 DCU上实现INT8 QK MMAC：分页访存、Fragment映射与GQA适配

gfx936 DCU上实现INT8 QK MMAC分页访存、Fragment映射与GQA适配前言本文是系列第二篇。第一篇《gfx936 DCU上实现INT8 KV与INT8 MMAC Attention推理优化》介绍了完整数据流本文聚焦 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.cuDecode 使用paged_attention_int8_kv带历史的分块 Prefill 使用gfx936_int8_prefix_attention_kernelPython 入口位于vllm/v1/attention/backends/rocm_aiter_fa.py。本文代码仓库https://github.com/sinpeyw/qwen3.5-vllm-dcu-optimization1. INT8 QK的量化和Scale恢复Attention Score 为S Q × K^T / sqrt(d)对 Query 和 Key 分别做对称动态量化Q ≈ Qq × sQ K ≈ Kq × sK代回原式可得S ≈ MMAC_INT8(Qq, Kq^T) × sQ × sK / sqrt(d)Qq、Kq是 INT8MMAC 输出是 INT32 accumulator。转成 FP32 后每个 score 需要乘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 MMAC2.1 Q放在operand AK放在operand B本文使用的整数矩阵指令为v_mmac_i32_16x16x32_i8一个 wave64 协同完成A: 16 × 32, INT8 B: 32 × 16, INT8 C: 16 × 16, INT32在这个布局中operand A 对应输出矩阵的行operand B 对应输出矩阵的列。QK 的自然映射是A Q输出行对应 Query row B K输出列对应历史 token我们早期也考虑过执行MMAC(K,Q)再转置。数学上两者等价但这种顺序会让输出归属、Key Scale 读取和 softmax 写入都变复杂。最终 CPU fragment 模型和 GPU 内核都固定为 Q 在 A、K 在 B。对于 Qwen3.5 的head_dim256INT8 指令每次推进 K 维 32因此只需256 / 32 8 个 K-step相同输出块的 BF16 指令每次推进 K 维 16需要 16 个 K-step。这个差异只属于矩阵乘加部分分页寻址、Scale 恢复和 softmax 不会同时减半。INT32 累加也不存在溢出压力。量化范围为[-127,127]时完整 256 维点积的绝对上界为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 的四个输出并不是四个连续列。令lane_col lane_id % 16一个 lane 实际持有的输出列为lane_col {0, 4, 8, 12}例如lane_col2时对应的是第2、6、10、14列而不是第2、3、4、5列。若共享内存写入、Key Scale 读取或尾部 mask 仍按连续四个 token 处理内核可以正常结束结果却会悄悄错位。QK 主循环可以简化为int32x4 acc{0,0,0,0};#pragmaunrollfor(intk_base0;k_base256;k_base32){int8x8 q_fragload_q_fragment(q_int8,k_base,lane);int8x8 k_fragload_k_fragment(k_cache,page,slot,k_base,lane);accmmac_i32_16x16x32_i8(q_frag,k_frag,acc);}intlane_collane_id%16;for(inti0;i4;i){intcollane_col4*i;intkey_tokentoken_basecol;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 前后保留等待asmvolatile(s_nop 4\n\tv_mmac_i32_16x16x32_i8 %0, %1, %2, %0\n\ts_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 如何消费它K cache: [block, kv_head, head_dim / 16, block_size, 16]加载一个 K fragment 时内核依次完成根据逻辑 key token 查询 block table计算该 token 在物理页中的 slot根据lane_k_base/16选择 head-dimension 小块用 8-byte 或 16-byte 向量 load 取得 INT8 fragment将 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 需要计算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量化时间40960.450 ms81920.896 ms163841.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 Decode16行中只有6种独立QueryDecode 每次只有一个新 token无法用 token 维填满 16 行。最终wxp内核按照 Query 头编号对 6 取模将 6 个 Query 头重复填入 16 个 MMAC 行。16 行都会执行但只有 6 行包含彼此不同的 Query因此独立有效行比例为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 tokenKV 头通过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 维从哪里取得并行度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 QK0.71x-1.00x转换和Scale抵消读取收益Query和Key直接执行INT8 MMAC原生BF16 QK大tile为1.55x-1.63x长上下文QK有明确收益原生INT8 MMAC小形状原生BF16 QK256-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 误差也可能改变最终概率排序。我们采用了四层验证CPU fragment模型重建 wave64 的16×16×32输入和输出确认每个 lane 的四个 accumulator 对应哪些列完整head_dim256INT8 对照覆盖 8 个 K-stepBF16 对照覆盖 16 个 K-step避免只计算一半维度分页和Scale覆盖跨物理 block、非零 slot、per-token Key Scale、尾部 mask 和 GQA head mapping运行时路径检查目标原生扩展、路由命中与真实 kernel 调用而不是只看二进制中存在 MMAC 指令。其中第二项来自一次真实教训早期 BF16 控制组只加载了 128 维 Query导致我们错误地把控制组问题归因成 INT8 误差。第三项也不能只验证连续 token因为 MMAC 输出列本身是交织的。最终这条专用路径明确检查head_size256、GQA ratio为6、无 sliding window和无 shuffle cachePrefill 路径还要求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 MMACV Scale融合、概率量化与Fragment分组》将继续介绍更复杂的 INT8 PV。