gfx936 DCU上实现INT8 PV MMAC：V Scale融合、概率量化与Fragment分组

gfx936 DCU上实现INT8 PV MMAC：V Scale融合、概率量化与Fragment分组

前言

本文是系列第三篇。第一篇介绍完整 INT8 KV 方案，第二篇介绍 INT8 QK；本文聚焦 Attention 的第二次矩阵乘法 Softmax(Score)×V，也就是 PV。

QK 的低精度公式比较直接：INT8 点积结束后，乘 Query Scale 和 Key Scale即可。PV更麻烦，因为每个历史 token 的Value都有自己的Scale，而PV恰好沿历史 token求和。这个Scale不能在整次矩阵乘法结束后统一补乘，必须先进入概率侧。

PV也是整条路线中最容易写错的一段。第一版内核运行很快，但接入服务后，两条正常回答都生成到了1024 token上限。我们最初怀疑 INT8 精度不足，逐lane重建fragment后才发现，问题不是量化公式，而是概率Scale的分组没有对齐 gfx936 MMAC实际消费的token。

修复Scale分组后，4K、8K、16K历史长度分别得到1.1049x、1.0979x、1.2653x，最大绝对误差不高于0.000977。这些数据只证明PV单项低精度计算成立，完整INT8链路的端到端与官方结果已经在第一篇给出。

对应代码位于csrc/rocm/attention.cu的INT8 PagedAttention和gfx936_int8_prefix_attention_kernel中。tools/xiandao_int8_pv_reference_check.py使用CPU fragment模型核对lane映射，csrc/rocm/int8_kv_cache.cu负责在上游生成INT8 V和对应Scale。

1. 为什么PV不能直接套用QK公式

对一个Query row，PV可以写成：

O = sum_i P_i × V_i

其中i遍历历史token，P_i是softmax概率。INT8 KV Cache为每个Value vector保存独立Scale：

V_i ≈ Vq_i × sV_i

代入可得：

O ≈ sum_i P_i × Vq_i × sV_i

sV_i随求和下标i变化，因此不能提出求和号。下面这种写法是错误的：

O ≈ MMAC_INT8(Pq,Vq) × 一个统一的sV

最终采用的变换是先把Value Scale融入概率：

W_i = P_i × sV_i
O ≈ sum_i W_i × Vq_i

然后对一个MMAC实际消费的32-token group中的W做动态量化：

sW = max(abs(W_group)) / 127
Wq = round(W / sW)

O_group ≈ MMAC_INT8(Wq,Vq) × sW

严格来说，送入MMAC的已经不是原始softmax概率，而是融合了per-token Value Scale的权重。文章仍将这段称为PV，是因为它对应Attention中的P×V阶段，但内核实现必须按W×Vq理解。

这套变换也能把误差来源拆开。记：

V_i = sV_i × Vq_i + eV_i
W_i = P_i × sV_i = sW_g × Wq_i + eW_i

则：

O – O_hat = sum_i P_i × eV_i
+ sum_i eW_i × Vq_i

第一项来自V Cache量化，第二项来自概率侧的分组量化。减小W的量化组可以缩小Scale覆盖范围，却会增加absmax归约和Scale数量。最终选择32 token，不只是为了方便，而是因为它正好对应一次INT8 MMAC的K维，能够避免额外的跨组归约。

2. 在online softmax中生成低精度概率

若先物化完整Score或概率矩阵，再启动独立量化kernel，会产生大量HBM读写。最终内核沿用FlashAttention的online softmax思路，逐个key tile处理。QK产生的INT32 accumulator先恢复为FP32 score：

score_i = int32_acc_i × sQ × sK_i / sqrt(d)

带历史Prefill和Decode都不会把完整概率矩阵写入全局显存，但两者进入PV前的权重形式并不完全相同。

2.1 带历史Prefill：量化未归一化权重

带历史Prefill在同一个Attention内核中维护：

m = 当前最大score
l = exp(score-m)的累计和
o = 当前输出累计

设旧状态为(m,l,o)，新tile合并后的最大值为m'：

alpha = exp(m-m')
p_tilde_i = exp(score_i-m')

然后更新：

o = alpha × o + sum_i p_tilde_i × sV_i × Vq_i
l = alpha × l + sum_i p_tilde_i

全部tile结束后才计算：

O = o / l

这一阶段量化的是当前online softmax基准下的未归一化权重：

W_i = p_tilde_i × sV_i

2.2 Decode：先在分区内归一化

Decode按256 token划分历史分区。每个分区先求出自己的最大值m_g和指数和l_g：

P_i = exp(score_i-m_g) / l_g
W_i = P_i × sV_i

随后再按MMAC实际拥有的32-token group量化W并执行PV。分区输出、m_g和l_g被保留下来，最后通过log-sum-exp权重归并多个分区。

因此，更准确的说法是：Prefill量化p_tilde×sV，Decode量化分区内归一化后的P×sV。两条路径的归一化位置不同，但共同约束是必须先把per-token Value Scale折叠进概率侧，再进行INT8量化。

这批权重只在寄存器和LDS中存在。每个MMAC-owned group依次执行：

计算W的absmax
→ 生成Scale
→ 将W量化为INT8
→ 读取分页INT8 V
→ 执行INT8 PV MMAC
→ 将结果按Scale恢复并累加到FP32输出

这条数据流避免了四类额外开销：

• 完整概率矩阵的全局写入和读取；
• 完整BF16 V中间张量；
• 独立的概率量化kernel；
• 每个token单独执行浮点P×sV×V_int8。

但概率量化本身也不是免费操作。每个group新增一次max reduction、一次倒数、舍入以及一个Scale。只有INT8 MMAC和INT8 V读取节省的时间大于这些成本，PV低精度路径才有意义。

3. P和V如何映射到MMAC

第二篇已经介绍过v_mmac_i32_16x16x32_i8的通用fragment映射，这里只讨论PV与该指令的对应关系。 PV中可以把16个Query row看作M，32个历史token看作K，16个Value输出维度看作N：

Wq: 16 × 32
Vq: 32 × 16
Out: 16 × 16

operand A拥有输出行，PV的输出行对应Query row，因此自然映射为：

A = quantized(P × sV)
B = V_int8

早期曾考虑沿用其他架构的顺序执行MMAC(V,P)，然后在共享内存中转置。数学结果等价，但lane输出归属与最终输出维度不一致，还会增加一次复杂转置。最终GPU内核和CPU fragment模型都固定为P在A、V在B。

对于同样的32-token K维，BF16需要两条16×16×16指令，INT8只需要一条16×16×32指令。同时，V不需要在寄存器中转成BF16，概率LDS的元素宽度也从16 bit降到8 bit。

核心数据流可以简化为：

// 当前online softmax基准下的未归一化权重，并已乘入V Scale。
float w[32] = exp_score_times_v_scale(...);

float amax = group_absmax(w);
float scale = amax / 127.0f;
int8x32 wq = quantize_i8(w,scale);

int8x32 vq = load_v_fragment_from_paged_cache(...);
int32x4 acc = mmac_i32_16x16x32_i8(wq,vq,0);

for (int i = 0; i < 4; ++i)
output[i] += float(acc[i]) * scale;

真实实现由64个lane共同完成，wq和vq分散在wave寄存器中，跨tile的FP32输出累加器还需要按alpha重标定。单个K=32 MMAC的INT32绝对上界为：

32 × 127 × 127 = 516,128

每个tile结束后，INT32结果立即乘回Scale并累加到FP32，因此整数累加器不需要跨长序列保存。

V Cache的布局也服务于这一方向：

V cache: [block,kv_head,head_dim,block_size]

token位于内层，一个lane可以为固定输出维度连续读取多个token的INT8 V。最终路径使用16-byte vector load；曾经尝试的8-byte load只装入了lane所需数据的一半，属于错误实现，而不是更轻量的性能候选。

4. Scale分组必须服从MMAC Fragment

这次问题定位花费了很长时间，也是PV最值得记录的一次错误。Decode PagedAttention中，一个wave覆盖64个token，需要两次K=32的INT8 MMAC。我们最初按逻辑连续区间计算概率Scale：

group 0: token 0-31
group 1: token 32-63

这在普通线性循环中很自然，却不是gfx936 fragment的实际消费顺序。两个MMAC真正取得的token分别为：

MMAC 0: {0-7,16-23,32-39,48-55}
MMAC 1: {8-15,24-31,40-47,56-63}

也就是说，一个MMAC的32个K元素由四段8-token slice交织组成。旧实现使用连续0-31和32-63的最大值生成Scale，却把结果应用到交织fragment，一条MMAC中会有一半概率使用错误的Scale。最终内核不再根据逻辑token区间猜测分组，而是直接根据MMAC输出归属计算：

const int token_offset = rowid + i * 4;
const int pv_tile = token_offset / 8;

rowid是16-lane子组编号，i是每个lane持有的四个输出。absmax、inverse Scale、概率写入和后续fragment load全部使用同一个pv_tile定义，不再维护两套容易分叉的分组规则。修复前，单独的量化参考误差并不大，native kernel却明显偏离BF16：

参考INT8 vs BF16: max_abs = 0.001293
native INT8 vs BF16: max_abs = 0.062622

这组差异说明，量化公式本身没有出现数量级异常，错误发生在硬件fragment映射。服务层的表现更明显。正常两条输出长度为：

[95,12]

错误PV内核却生成：

[1024,1024]

表面输出吞吐从7.996790 tok/s升到19.200604 tok/s，TTFT和TPOT却没有出现对应幅度的改善。所谓“2.4倍吞吐”来自模型一直不能生成EOS，不是性能提升。

从这次问题以后，PV验证必须同时满足：

• Scale按MMAC实际拥有的token计算；
• Wq写入和fragment读取使用同一映射；
• CPU reference使用同一token group；
• native-vs-reference和quantized-vs-BF16同时通过；
• 服务输出长度和文本保持合理。

只检查kernel是否完成、输出是否为有限数，无法发现这类错误。

5. 修复后的PV算子结果

修复Scale分组后，我们在同一个原生INT8 Cache Attention内核中，仅切换BF16 PV与INT8 PV：

上下文越长，INT8 PV的收益越明显。固定的概率量化、Scale和同步开销会被更多V读取与矩阵乘法摊薄。

这张表的对照边界主要注意：两侧已经共同使用INT8 Cache和INT8 QK，唯一变量是PV内部采用BF16矩阵计算还是INT8 MMAC。因此它能证明“PV改用INT8 MMAC”在该内核中有收益，但还不能证明从完整BF16服务切换到完整INT8服务会得到同样比例。

最终INT8 PV与INT8 QK、融合producer和选择性Prefill一起进入正式提交。完整链路的服务数据和官方评测见第一篇，平台没有提供PV单项拆分。

6. 踩坑经历

除了前面的Scale分组错误，我们还在访存宽度、分区大小和量化粒度上走过几条弯路。这些尝试有的性能不理想，有的则是看起来更快，实际少计算了数据。

最典型的是8-byte V load。它看起来比16-byte load更灵活，实测速度也一度更快，但逐lane核对后发现，它只取得了当前fragment所需16个token的一半。这不是一个较快的优化版本，而是没有完成全部计算的错误实现。最终路径恢复为16-byte vector load。类似的问题也出现在partition 512上：早期实现的地址映射与four-wave launch不一致，会越过每个fragment实际拥有的数据范围；即使修正寻址，性能也没有稳定优势，因此Decode最终保留256-token partition。

我们还尝试过为整个partition只计算一个probability Scale，希望减少Scale数量。但这种做法需要额外的跨wave max reduction，增加同步开销，而且更大的量化范围会降低有效精度。最终选择让每个MMAC实际拥有的32-token group独立计算Scale。这个粒度既与16×16×32指令的K维一致，也不需要额外跨组归约。

另一条看似稳妥的路线是先把INT8 V恢复为完整BF16 V，再调用原有PV。这样虽然容易复用成熟内核，却会增加中间张量的写入和读取，也完全用不上INT8 MMAC。即使不写回全局显存，只在寄存器中逐fragment反量化，类型转换和Scale指令仍可能抵消INT8 V减少的读取开销。

概率采用UINT8也曾被列为候选。softmax概率和P×sV都非负，UINT8的0-255比signed INT8的0-127多一位有效精度，理论上具有吸引力。但V使用对称signed INT8，概率单独改为UINT8还需要可用且高效的mixed-signed指令路径，或者对V增加偏移和补偿；后者会引入新的归约与校正。比赛结束前，我们没有得到一条同时通过正确性和性能验证的UINT8×INT8路径，因此最终代码只保留signed INT8。当前方案按每个32-token group独立计算Scale，误差已经控制在1e-3量级；对实际内核而言，减少转换、补偿和同步，比单纯增加一个bit的正数编码范围更重要。

7. 结语

将Softmax×V改成INT8 MMAC，核心不是把概率直接转换为INT8，而是先把per-token Value Scale正确折叠进online softmax权重，再按照MMAC实际消费的32-token fragment计算Scale和量化。

最终可以确认：修复fragment分组后，4K、8K、16K的PV单项分别达到1.1049x、1.0979x、1.2653x，最大绝对误差不高于0.000977。

QK的困难主要是分页K供数和GQA映射，PV的困难则是Scale位于求和维度、概率需要再次量化，以及数学分组必须服从硬件fragment。两部分共同成立，INT8 KV才真正从存储格式变成可被低精度矩阵单元直接消费的数据。