# gfx936 DCU上实现INT8 PV MMAC：V Scale融合、概率量化与Fragment分组 ## 前言 本文是系列第三篇。[第一篇](https://blog.csdn.net/Foreverahead/article/details/162978554?spm=1001.2014.3001.5501)介绍完整 INT8 KV 方案，[第二篇](https://blog.csdn.net/Foreverahead/article/details/163175989?spm=1001.2014.3001.5501)介绍 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。 本文代码仓库：[https://github.com/sinpeyw/qwen3.5-vllm-dcu-optimization](https://github.com/sinpeyw/qwen3.5-vllm-dcu-optimization) ## 1. 为什么PV不能直接套用QK公式 对一个Query row，PV可以写成： ```text O = sum_i P_i × V_i ``` 其中 `i`遍历历史token，`P_i`是softmax概率。INT8 KV Cache为每个Value vector保存独立Scale： ```text V_i ≈ Vq_i × sV_i ``` 代入可得： ```text O ≈ sum_i P_i × Vq_i × sV_i ``` `sV_i`随求和下标 `i`变化，因此不能提出求和号。下面这种写法是错误的： ```text O ≈ MMAC_INT8(Pq,Vq) × 一个统一的sV ``` 最终采用的变换是先把Value Scale融入概率： ```text W_i = P_i × sV_i O ≈ sum_i W_i × Vq_i ``` 然后对一个MMAC实际消费的32-token group中的 `W`做动态量化： ```text 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`理解。 这套变换也能把误差来源拆开。记： ```text V_i = sV_i × Vq_i + eV_i W_i = P_i × sV_i = sW_g × Wq_i + eW_i ``` 则： ```text 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： ```text score_i = int32_acc_i × sQ × sK_i / sqrt(d) ``` 带历史Prefill和Decode都不会把完整概率矩阵写入全局显存，但两者进入PV前的权重形式并不完全相同。 ### 2.1 带历史Prefill：量化未归一化权重 带历史Prefill在同一个Attention内核中维护： ```text m = 当前最大score l = exp(score-m)的累计和 o = 当前输出累计 ``` 设旧状态为 `(m,l,o)`，新tile合并后的最大值为 `m'`： ```text alpha = exp(m-m') p_tilde_i = exp(score_i-m') ``` 然后更新： ```text o = alpha × o + sum_i p_tilde_i × sV_i × Vq_i l = alpha × l + sum_i p_tilde_i ``` 全部tile结束后才计算： ```text O = o / l ``` 这一阶段量化的是当前online softmax基准下的未归一化权重： ```text W_i = p_tilde_i × sV_i ``` ### 2.2 Decode：先在分区内归一化 Decode按256 token划分历史分区。每个分区先求出自己的最大值 `m_g`和指数和 `l_g`： ```text 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依次执行： ```text 计算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： ```text Wq: 16 × 32 Vq: 32 × 16 Out: 16 × 16 ``` operand A拥有输出行，PV的输出行对应Query row，因此自然映射为： ```text 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。 核心数据流可以简化为： ```cpp // 当前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绝对上界为： ```text 32 × 127 × 127 = 516,128 ``` 每个tile结束后，INT32结果立即乘回Scale并累加到FP32，因此整数累加器不需要跨长序列保存。 V Cache的布局也服务于这一方向： ```text 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： ```text group 0: token 0-31 group 1: token 32-63 ``` 这在普通线性循环中很自然，却不是gfx936 fragment的实际消费顺序。两个MMAC真正取得的token分别为： ```text 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输出归属计算： ```cpp 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： ```text 参考INT8 vs BF16: max_abs = 0.001293 native INT8 vs BF16: max_abs = 0.062622 ``` 这组差异说明，量化公式本身没有出现数量级异常，错误发生在硬件fragment映射。服务层的表现更明显。正常两条输出长度为： ```text [95,12] ``` 错误PV内核却生成： ```text [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： | 历史长度 | BF16 PV | INT8 PV | 加速 | 最大绝对误差 | RMS误差 | | -------: | -------------: | ------------: | ----------: | -----------: | -----------: | | 4K | `57.149 us` | `51.725 us` | `1.1049x` | `0.000977` | `0.000252` | | 8K | `74.186 us` | `67.568 us` | `1.0979x` | `0.000732` | `0.000179` | | 16K | `120.989 us` | `95.619 us` | `1.2653x` | `0.000488` | `0.000122` | 上下文越长，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才真正从存储格式变成可被低精度矩阵单元直接消费的数据。