RDNA3上llama.cpp的Flash Attention：比Vulkan f16 K减少47% KV VRAM，在F16 K / q4_0 V上KLD几乎无损。第一部分。

Reddit r/LocalLLaMA 工具
llama.cpp rdna3 flash-attention kv-cache optimization memory amd

摘要

一种针对RDNA3 GPU上llama.cpp的新packed16 K技术，相比Vulkan fp16将KV缓存VRAM减少47%，使用int8打包和原生dot4指令，以最小的KLD损失保持fp16质量的K值。

llama.cpp注意力机制中的常规权衡是：量化KV缓存会损失质量，或者保留fp16会消耗大量VRAM。在RDNA3上，从此有了第三种选择！将四个8位K值打包成一个32位值，直接馈入GPU的原生`sudot4`点积指令。没有K的有损量化。没有fp16 K缓冲区驻留在内存中。内核获得它所需的确切数据布局，VRAM下降，因为存储的是8位K负载加上fp16缩放因子，而不是完整的fp16 K张量。但真正的差异出现在128k上下文且运行活动MTP草稿模型时——现在你要为*两个*完整上下文（主模型+草稿模型）存储K和V。通过`rocm-smi`测量的总VRAM：128k活动MTP，两侧q4_0 V | | Vulkan f16 K | |--------|--------------| | 23.18 GiB | 22.50 GiB | | ROCm packed16 K** | **21.76 GiB** | 这1.42 GiB的差异决定了能否容纳128k MTP会话，具体取决于其他VRAM压力。这不是模型权重的节省（这些权重是相同的）——纯粹是削减了两个上下文的K缓存内存占用。现在来看质量方面。packed16 K路径在反量化后仍然产生fp16范围的K值——8位打包不是有损量化，而是存储布局的改变。唯一的压缩损失来自V侧。在WikiText-2上使用27B模型进行测量，ctx=512，chunks=4，将V=q4_0和V=q8_0与V=fp16基线进行比较。所有候选中K均为packed16 I32： **q4_0 V vs fp16 V:** | 指标 | 值 | |------|------| | Mean PPL ratio | 1.0020 ± 0.0042 | | Mean KLD | **0.00455** ± 0.00034 | | Median KLD | **0.00182** | | 99th percentile KLD | 0.0500 | | Same top token | **97.06%** | | RMS Δp | 1.98% | **q8_0 V vs fp16 V:** | 指标 | 值 | |------|------| | Mean PPL ratio | 1.0010 ± 0.0034 | | Mean KLD | **0.00283** ± 0.00033 | | Median KLD | **0.00086** | | 99th percentile KLD | 0.0313 | | Same top token | **97.94%** | | RMS Δp | 1.68% | 关于这些KLD数值的含义：Kullback-Leibler散度衡量两个概率分布之间的差异程度。在~0.01以下通常被认为在实际的token级别分布中几乎无法区分。两种V格式都远低于该阈值，其中q8_0的散度大约是q4_0的一半（均值0.0028 vs 0.0046，中位数0.0009 vs 0.0018）。如果你使用q4_0 V以保持精简，你需要为比fp16 K+V更少的KV VRAM付出~0.0045 KLD。如果你想要更严格的质量，q8_0 V相比fp16 K+V给出~0.0028 KLD（因为K的节省是相同的，V格式不会改变packed16 K布局）。为什么packed16 K能产生等效于fp16的质量？因为打包不是量化，而是重新打包。K张量在空闲时是fp16格式。内核读取每一行，计算每个块的fp16缩放因子（absmax），即时量化为int8，将四个int8值打包成一个I32，然后将该负载和缩放因子写入缓存。在注意力传递中，内核加载I32负载，调用`sudot4`（该指令在一个指令中执行四个INT8乘法和一个累加），乘以Q和K的缩放因子，然后进行在线softmax。对于打包后的int8范围，反量化在数学上是精确的！唯一的信息损失是K值的int8舍入，而这受每个块的fp16缩放因子限制。WikiText的数字证实了这一点：1.002的PPL比例完全在±0.004的噪声带内。与Vulkan的做法相比：在Vulkan上，KV缓存路径将K存储为完整的fp16。这对K是无损的，但消耗内存。packed16方法能给你相同有效K精度（int8舍入加fp16缩放因子实际上等于fp16范围），同时将K内存占用削减到大约三分之一——每个值8位加上缩放开销，对比16位。V侧也减半。对于有效的4_0 V，你得到2.25位。 [https://github.com/DrBearJew/llama.cpp/tree/tbq4-rdna3-experiment](https://github.com/DrBearJew/llama.cpp/tree/tbq4-rdna3-experiment) [https://github.com/DrBearJew/dot4-flash-attention](https://github.com/DrBearJew/dot4-flash-attention)
查看原文

相似文章

双GPU llama.cpp加速

Reddit r/LocalLLaMA

llama.cpp的一个分支修复了量化KV缓存中的--split-mode tensor问题，在双GPU配置上实现高达40%的速度提升，且无质量损失。