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

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。

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};

#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 前后保留等待：

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

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更适合：

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分组》将继续介绍更复杂的 INT8 PV。