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.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

1. 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是 INT8，MMAC 输出是 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 MMAC

2.1 Q放在operand A，K放在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_dim=256，INT8 指令每次推进 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_col=2时，对应的是第2、6、10、14列，而不是第2、3、4、5列。若共享内存写入、Key Scale 读取或尾部 mask 仍按连续四个 token 处理，内核可以正常结束，结果却会悄悄错位。

QK 主循环可以简化为：

int32x4 acc={0,0,0,0};#pragmaunrollfor(intk_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);}intlane_col=lane_id%16;for(inti=0;i<4;++i){intcol=lane_col+4*i;intkey_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 前后保留等待：

asmvolatile("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 如何消费它：

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量化时间
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，因此独立有效行比例为：

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 维从哪里取得并行度：

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 误差也可能改变最终概率排序。我们采用了四层验证：

CPU fragment模型：重建 wave64 的16×16×32输入和输出，确认每个 lane 的四个 accumulator 对应哪些列；
完整head_dim=256：INT8 对照覆盖 8 个 K-step，BF16 对照覆盖 16 个 K-step，避免只计算一半维度；
分页和Scale：覆盖跨物理 block、非零 slot、per-token Key Scale、尾部 mask 和 GQA head mapping；
运行时路径：检查目标原生扩展、路由命中与真实 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更适合：