如果你最近在跑长上下文模型大概率会撞上同一个墙模型权重已经量化到 INT4 了显存还是不够用。我自己踩的坑是在一块 48G 的卡上跑 128K 上下文的业务模型FP16 权重占 14GKV Cache 却吃掉了 64G 还要多——权重量化省下来的那点空间根本追不上 KV Cache 的增长速度。这也是为什么看到 MXAttention 加 MXFP4 这个组合时我专门花几天把格式规范和相关实现从头到尾过了一遍。这篇文章就当是我整理后的笔记聊聊 MXFP4 这种 4 位浮点格式到底解决了什么痛点MXAttention 又是怎么把它用到注意力机制的 KV Cache 压缩上以及你如果要部署需要避开哪些坑。1. 长上下文推理的显存黑洞权重量化解决不了的事1.1 KV Cache 是怎么一步步吃光显存的先说清楚 KV Cache 的账怎么算。自回归模型每生成一个新 token都要计算当前 token 对应的 Key 和 Value 向量然后跟历史上所有 token 的 Key、Value 拼在一起做注意力计算。这些历史 K、V 不会凭空消失而是被缓存下来供后续每个 token 复用这就是 KV Cache。它的显存占用公式很简单KV 字节数 2K 和 V 两份× 层数 × KV head 数 × head 维度 × 序列长度 × 每个元素字节数 × batch size。拿一个常见的 7B 级模型举例32 层、32 个 KV head、head dim 128在 FP16 下每个 token 每层要存 16KB 的 KV32 层合计就是 512KB。跑 128K 上下文时KV Cache 总量能到 64GB。是的你没看错权重只有 14GB 的模型KV Cache 是权重的 4 倍还多。这时候哪怕权重压到 INT4也只能腾出 10GB 出头KV Cache 依然放不下。而且 decode 阶段每个新 token 都要把这 64GB 的 KV 从头到尾读一遍做注意力计算。这带来两个问题一是容量不够就会换出到 CPU 内存生成速度断崖式下跌二是就算容量够带宽也会被 K/V 的反复读取吃满。长上下文推理的瓶颈早就不是权重的容量和算力而是 KV Cache 的容量和带宽。1.2 权重能压KV 为什么压不动有人会问权重量化都这么成熟了KV Cache 照搬不行吗这里有个很反直觉的地方权重是离线已知的静态张量你可以在训练后拿一批校准数据统计激活分布做全局的 scale 选择和精度补偿。但 KV Cache 是推理过程中逐 token 在线产生的你根本不知道未来会出现什么样的值没有机会做全局校准。更麻烦的是 K 和 V 的分布结构很不一样。K 向量经常出现明显的 outlier 维度某些特征通道的数值会比其他通道高出好几个数量级V 的分布相对平缓更像输入激活的平滑版本。同一个量化策略很难同时照顾好这两个分布完全不同的张量。还有误差传播的问题。权重量化错了影响是静态的最多是推理结果略微漂移。KV Cache 量化错了误差会随着生成过程逐 token 累积早期 token 的量化噪声改变了注意力权重影响了后续 token 的 K/V进而让后面的误差越滚越大。这就是为什么很多人在 KV 上做低比特量化时指标看着还行一跑长文本就暴露幻觉、重复、上下文丢失。1.3 既有 KV 压缩方案的共同天花板目前主流方案大致分三类。第一类是结构压缩比如 GQA分组查询注意力和 MQA通过让多个 query head 共享同一组 K/V head 来减少 KV 的份数。这个方法有效但它是模型结构层面的改了结构就得改模型已经训好的模型用不上。第二类是 INT8/INT4 的 KV 量化。这类方案的问题在于整数格式的数值网格是固定刻度的INT4 只有 16 个均匀分布的档位必须靠一个全局或 per-token 的 scale 把整个张量压到 [-8, 7] 区间。一旦某个 block 里出现一个 outlier要么把 scale 拉大导致所有正常值精度坍缩要么把这个 outlier 单独剔除增加额外处理逻辑。第三类是各种混合精度策略比如95% 的 KV 用低比特、少数 outlier 通道用高比特的方案。这类方法精度不错但需要额外的 outlier 检测表和动态路由在解码阶段会引入分支判断对 GPU kernel 不友好实际加速并不理想。这三类方案的天花板其实就是一句话整数格式的动态范围太窄而 KV 的在线生成又注定了没法靠全局统计来兜底。想要在 4 bit 这种位宽下同时保住动态范围和局部自适应能力就得换成一种更聪明的数值格式。2. MXFP4 不是又一个 INT4微观缩放格式的底层逻辑2.1 从 INT8 到 FP8 再到 FP4数值格式演进的动机先回顾一下浮点格式为什么会在这个阶段重出江湖。FP8 流行起来之后大家发现对于大模型推理动态范围比均匀精度更重要。INT8 在 [-127, 127] 之间只有 256 个均匀档位表达大型激活值没问题但小数部分精度很差FP8 的 E4M3 和 E5M2 因为有指数位能把很大和很小的数都表达出来所以对 outlier 的天然耐受度更高。顺着这条思路往 4 bit 走问题就来了FP4 如果把格式砍成 1 位符号 2 位指数 1 位尾数E2M1尾数只有 0.5 的步长数值集合极其稀疏如果把尾数留宽砍指数动态范围又没了。硬从 8 位砍一半是不可行的必须引入额外的机制来弥补。MX 规范Microscaling Formats由 OCP 主导的一套开放数值格式标准给出的答案是4 位只存不完整但高动态的尾数和指数另加一个共享的微观缩放因子microscaling factor让一小块元素共享这个 scale。MXFP4 就是这套规范里的 4 位浮点变体。它不只是又一种低比特格式而是用 block-level scale 把有限位宽和全局动态范围这两个矛盾解开了。2.2 E2M1 的数值结构与动态范围MXFP4 在存储上每元素只有 4 bit布局是 1 位符号、2 位指数、1 位尾数。你可以把它理解为二进制的科学计数法尾数只提供 1.0 和 1.5 两个有效数字隐式前导 1 加一位尾数但指数有 4 种刻度决定了数值落在哪个数量级。单独看这 4 个 bit它的数值集合确实稀疏——同一个指数刻度下只有两个尾数值一共十几个有效数字。但 MXFP4 的设计意图从来不是独立的 4 位浮点而是永远搭配一个 scale 使用。块内所有元素共享一个外部缩放因子相当于把整个稀疏数值集合在数轴上平移和伸缩。加上这个 scale 之后它实际能表达的范围横跨非常宽的数量级对极端 outlier 有天然的抗性。有个很直观的类比INT4 就像一把固定刻度的尺子量程只有 15 格放一个 3.7 的数进去精度就只有 0.5 左右MXFP4 不是尺子是一个可调倍率的显微镜每个小区域都可以先自动对焦到合适的倍数再去做那一格一格的量化。同样是 4 bit前者的量化误差取决于全局值的分散程度后者只取决于局部块的分散程度。2.3 块级共享缩放因子量化精度的关键MXAttention 使用的 MXFP4 量化过程核心就是每块元素共享一个 scale。具体说一段连续 K/V 向量会被切成若干 block每个 block 通常 16 或 32 个元素。对每个 block 单独执行找出块内绝对值的最大值记为 amax根据 amax 计算一个 2 的幂次作为 scale比如 scale 2^ceil(log2(amax))这样后续除法在硬件上就是一次指数减法或移位成本极低块内每个元素除以 scale映射到 E2M1 能表达的范围内做四舍五入存成 4 bit 的指数尾数编码scale 单独用 FP8 或指数形式存下来每个 block 一份。反量化就更简单了读回 4 bit 编码和它所属 block 的 scale两者相乘就是近似原始值。整个过程没有复杂的查找表没有昂贵的除法只有移位、比较和乘法正好是 GPU kernel 最擅长的操作。有一个细节值得注意scale 的选择不需要非常精确用 2 的幂次近似就够了。原因是 MXFP4 本身有指数位scale 差一个 2 的因子只是让块内所有数值统一平移一个刻度块内元素之间的相对关系不会变。这种宽容度让硬件实现变得非常干净。2.4 与 INT4 的对比从两级坍缩到局部自适应用过 INT4 做 KV 量化的朋友应该体会过两级坍缩第一级是张量级的 scale 被 outlier 带偏导致大部分正常值的有效步长变得很粗第二级是即便用了 per-token 或 per-channel scale一个 token 内部如果同时存在大数和小数小数部分仍然会被碾碎。MXFP4 的 block 机制天然规避了这个问题。每个 block 的 scale 只看自己块内的最大值块与块之间互不干扰。一个包含 outlier 的 blockscale 会变大块内其他元素确实会损失一些精度但旁边那些没有 outlier 的 blockscale 完全不受影响依然保持着精细的相对精度。换句话说出问题的代价被局部化了不会像全局 scale 那样一颗老鼠屎坏了一锅汤。下面这张表是我整理的核心对比对比维度INT4MXFP4E2M1 block scale存储位宽4 bit4 bit 每 block 一个 scale数值布局均匀网格16 个档位科学计数法指数刻度 块级缩放有效动态范围很窄靠全局 scale 硬撑很宽块级 scale 自动适配对 outlier 的抵抗力低outlier 会拉垮全块高outlier 只影响所在 block是否需要全局校准通常需要不需要每块在线算 max 即可硬件成本成熟 INT 指令存 4 算 16 可绕开原生支持正在普及3. MXAttention 的设计把 MXFP4 塞进注意力机制3.1 量化单元的选择为什么必须 per-block把 MXFP4 用到 KV Cache 上第一个要定的事是 block size。理论上有三种口径per-tensor整层一个 scale、per-token每个 token 一套 scale、per-block每 16/32 个元素一套 scale。per-tensor 基本不用考虑和 INT4 撞同样的问题。per-token 比 per-tensor 好但它依然没解决同一 token 内 K 向量不同维度的数值差异——KV Cache 里每个 token 是一个向量序列维度之间经常有数量级差一个 scale 覆盖一整条向量还是太粗。per-block 是存储开销与精度收益的平衡点。K/V 的 head dim 通常是 64 或 128block 取 16 或 32 都能整除每个 block 内元素天然来自同一个特征子空间分布相对同构量化噪声最小。同时GPU 的向量访存和 Tensor Core 做矩阵乘法时本来就是一串连续元素一起处理的block 对齐就和硬件执行粒度完美重合。我自己试下来的经验head dim 是 128 时优先试 block32精度基本能接受如果 PPL 掉了超过预期再降回 16。block 越小精度越稳但 scale 的存储开销会增加这个后面算账。3.2 K 和 V 的分开处理MXAttention 里有个很容易被忽略的细节K 和 V 不应该用完全相同的量化参数。原因回到第一节说的分布差异上。Value 向量在注意力计算中最终要跟注意力权重做加权求和它的分布更像模型中间激活的平滑分布绝对值跨度较小用 MXFP4 量化块内相对精度保持得不错。Key 向量则有更明显的方向性某些 head 的 K 会出现持续较高的激活值很像特征检测器outlier 更频繁。好在 MXFP4 的 block scale 会自动把 outlier 单独兜住不需要额外分离通道。实操时我建议先只量化 V、K 保持 FP16 或者 FP8跑一遍验证集。如果精度损失在可接受范围再尝试把 K 也降到 MXFP4。这样定位问题快不会被 K、V 同时引入的噪声搞晕。实际测试中V 量化到 MXFP4 的收益已经很大因为 V 的读取在 Flash Attention 的加权求和阶段是带宽大头。3.3 内存账本压缩率到底怎么算现在认真算一笔账。FP16 下 K/V 每个元素 16 bitMXFP4 下每个元素 4 bit额外每个 block 要存一个 scale。scale 常用 8 bit 浮点或 8 bit 指数表示。block16 时每元素平均 4 8/16 4.5 bit block32 时每元素平均 4 8/32 4.25 bit。换算成压缩率block16 约 3.56 倍block32 约 3.76 倍。如果再算上内存 pool 的对齐 padding 和元数据开销实际收益通常是 3 到 3.5 倍。回到开头那个 64GB 的例子KV Cache 用 MXFP4 之后大概降到 18-20GB。也就是说48G 的卡不只是刚好放下还能把 batch 翻一倍或者把上下文从 128K 提到 256K。这个数量级的收益权重量化给不了。可能有人会问scale 本身也占内存会不会白忙注意 scale 是每 16 或 32 个元素共享一个均摊下来每元素只多 0.25 到 0.5 bit相对从 16 bit 降到 4 bit 省下来的 11.5 bit完全不成比例。3.4 与 Flash Attention 融合的推理流程KV Cache 量化不是把数据存成 4 bit 就完事了关键在推理流程里怎么读、怎么算。预填充阶段prefill输入 prompt 一次处理完全部 K/V 计算出来之后立即按 block 做 MXFP4 量化写入低比特缓存。解码阶段decode注意力 kernel 按 block 读取 4 bit KV取回 scale 反量化成 FP16/BF16再做标准的 Flash Attention 计算。这个过程看似多了一步反量化但实际上很划算decode 是访存密集型操作K/V 从 16 bit 变成 4 bit内存带宽压力直接降了约 3.5 倍反量化的乘法和移位操作相对注意力矩阵乘法来说成本很低整体收益远大于开销。更进阶的做法是把反量化融合进 Flash Attention 内核在读取 K/V tile 时顺便用 scale 缩放避免额外遍历一遍数据。如果底层硬件支持原生 FP4 的矩阵乘比如新一代推理加速器开始支持 MX 格式那还可以直接让 Tensor Core 在 4 bit 精度下做 QK^T 和 PV 乘法把计算和访存双重降下来。4. 实测收益显存、吞吐与精度的真实取舍4.1 显存占用从 64G 到 20G 的账本我在三个不同量级的模型上复现过 MXAttention 思路的 KV 量化7B、14B、72B。前提条件是只对 KV Cache 做 MXFP4 量化权重保持原精度不做额外压缩block32K 和 V 都量化scale 用 8bit 指数存储。结果很有规律核心收益完全由 KV Cache 占比决定。7B 模型跑 128K 上下文时 KV Cache 占大头压缩后整体显存占用直接降了约 65%14B 模型权重占比略高整体降幅约 50%72B 模型如果只跑 32K 上下文KV 占总显存比例本身就低收益就缩水到 20% 左右。教训也很明显KV 量化确实是把显存从长上下文中解放出来的关键但前提是场景里 KV 本身是主导项。如果你本来就用不到长上下文或者 GQA 已经把 KV 压得很小那么上 MXFP4 的意义不大。4.2 吞吐与延迟收益decode 阶段每生成一个 token需要读取全部历史 KV。假设 batch16、上下文 64KKV Cache 60GB那么每生成一个 token 要读 60GB 数据在权威的 GPU 显存带宽 2TB/s 左右的情况下理想状态纯读就已经要花 30ms 级别的时间。把 KV 用 MXFP4 压缩到 4.5 bit 之后同样场景读入量降到约 17GB理论耗时降到 8ms 左右。这就是吞吐提升的来源。我的实测数据是在 batch 足够大、带宽吃满的配置下整体 decode 吞吐能提升 1.5 到 2 倍batch 很小的时候提升不明显因为单条请求的注意力计算还有反量化等额外开销。有一点要特别提醒如果你只是把 KV 存成 4 bit反量化却用了一个独立的 kernel 先回写 FP16那访存压力一点没减甚至可能更慢。MXAttention 的效果必须建立在反量化融合进注意力 kernel 的前提下单纯换数据格式不换 kernel属于只交学费不拿收益。4.3 精度表现与敏感层精度是我最担心的一块实际测下来比预期好。用困惑度PPL和摘要类下游任务做评测block32、K/V 全部 MXFP4 时PPL 通常只上升 0.1 到 0.3block16 几乎可以做到无感。但有个规律必须强调不是每一层都一样皮实。浅层前 1/3的注意力头对 KV 量化容忍度很高深层尤其是最后几层的 V 向量会更容易暴露精度问题。我建议先整层量化跑一遍评测如果 PPL 超标优先把最后 3 到 4 层退回 FP16 或者 block16而不是全量降低 block size。这样能把精度损失补回来同时只牺牲很小一部分显存收益。另外长上下文下要额外关注生成质量而不只是 PPL。PPL 是全局平均掩盖了局部崩坏。我在测试中发现有些模型在 64K 之后会开始出现重复片段KV 量化让这种现象提前了。所以长文本压力测试要有专门检查重复率和上下文召回的手段别只看 PPL。4.4 哪些场景值得上、哪些不值得我把自己的判断标准列成一个简单的筛选表场景是否推荐 MXAttention原因上下文 32K 以上KV 占显存过半强烈推荐省 60% 以上显存精度代价低高并发大 batch 推理推荐decode 带宽瓶颈明显吞吐提升大24G 消费级显卡跑长上下文推荐是少数能让 16K 上下文跑进显存的办法8K 以下短上下文不推荐收益小反量化开销反而占比重对单 token 精度极其敏感的用途谨慎建议只量化 V或按层回退还在训练/微调阶段的模型不推荐这是部署期优化不适合训练期5. 从原理到部署MXAttention 落地要过的几道关5.1 硬件支持的现状判断先说结论就算你的卡完全不支持 FP4 原生计算MXAttention 的 KV 压缩依然能用因为落地方式可以做成存 4 bit、算的时候反量化回 FP16。这条路对任何 GPU 都成立收益来自显存容量和带宽跟是否有 FP4 指令无关。想要进一步吃到4 bit 直接算的红利就需要硬件和推理框架支持 MX 格式了。当前新一代推理加速器普遍开始宣传对 FP4/MXFP4 的支持不同代际和 SDK 的暴露程度差别很大。我的建议是先查你正在用的推理框架有没有 KV cache 量化选项有的话直接开没有的话先从手写 kernel 或算子原型验证效果别等框架原生支持。5.2 量化算子的实现细节与精度坑真正写量化算子时有几个坑特别值得说。第一个是四舍五入模式。我最早用 floor 截断实现PPL 明显比 round-to-nearest-even 高出一截。MX 规范里默认推荐 round-to-nearest-even这个小细节对低比特格式的影响非常大因为你总共只有 4 个 bit任何系统性偏差都会直接变成精度损失。第二个是 scale 的计算。直接用块内 max 作为 scale在无 outlier 的 V 向量上没问题但 K 向量如果恰好某个维度出现异常大值一个 block 的精度会被整体拉低。更稳的做法是取块内百分位或绝对值均值乘个安全系数来定 scale牺牲极少动态范围换取有效精度。下面给一个用 NumPy 演示核心思想的简化版代码注意真正的工程实现要按 MX 规范处理 E2M1 编码和硬件指令我这里只展示 scale 计算与反量化的逻辑骨架import numpy as np def mx_block_quantize_dequantize(x, block_size16): # x: 一维连续向量长度会被看成 n_blocks * block_size n_blocks x.shape[0] // block_size x x[: n_blocks * block_size].reshape(n_blocks, block_size) # 1. 每块取绝对值最大值 amax np.abs(x).max(axis-1, keepdimsTrue) # 2. scale 取向上取整的 2 的幂方便硬件用移位实现 scale np.power(2.0, np.ceil(np.log2(np.maximum(amax, 1e-12)))) # 3. 元素除以 scale进入 E2M1 的表示区间 scaled x / scale # 4. 简化模拟 4bit 表示按 0.5 步长四舍五入并 clamp quantized np.round(scaled / 0.5) * 0.5 quantized np.clip(quantized, -1.5, 1.5) # 5. 反量化乘回 scale恢复近似原始值 dequantized quantized * scale return dequantized.reshape(-1)第三个是 block 与 head dim 的对齐。head dim 是 128 时block32 正好 4 块如果是 64 的 head dimblock16 更合适。设计 kernel 时提前把 block size 固定成编译期常量让循环展开和寄存器分配更友好性能差距比想象中大。5.3 长上下文压力测试与校准策略部署前务必跑三轮验证第一轮PPL 对比。用一段有代表性的长文本最好 16K 以上计算 PPL和 FP16 baseline 对比记录不同 block size、K/V 量化组合下的差值。差值 0.3 以内基本可用超过 0.5 就要调整策略。第二轮行为异常检测。量化的 KV 在长上下文中容易引发重复率和注意力涣散。准备一个长文本问答集看模型能不能准确引用上下文前部的事实同时统计输出中 n-gram 重复率对比 baseline。第三轮压力变化测试。把上下文从 16K 逐步拉到 64K、128K观察 PPL 是否在某个长度上突然劣化。这种现象往往指向量化误差在长距离注意力路径上累积解决办法就是按层回退具体做法前面说过优先保住最后几层的精度。5.4 工程落地阶段我最想提醒的几件事如果只能带走三句话我希望是这三句。第一句不要为了省几个 bit 把所有层一律打死。KV 量化最大的自由度就是你有很多层很多块可以分别选择精度这是一个配置组合问题不是一把梭的问题。第二句scale 存储不要随便塞在主 KV pool 里。我会把 scale 单独开一块显存池因为它的访问模式和 KV 数据不同KV 是每个 token 都要从头读到尾scale 是每个 block 读一次。混在一起会破坏访问模式导致 cache locality 变差。第三句版本质量评估里一定要有一条最长上下文 最高并发的组合用例。KV 量化的问题很多时候在小规模测试里根本暴露不出来只有在显存带宽被打满、上下文拉满的时候才会以生成质量劣化的形式冒出来。我看过太多团队在评测集上跑得好好的上生产就被用户反馈模型失忆。MXAttention 这个方向的本质是承认 KV Cache 已经成了长上下文推理的头号瓶颈然后用一个新数值格式把存储密度和动态范围同时提上去。它在显存、带宽、精度上的取舍比传统 INT 量化要优雅得多。如果你也在被长上下文显存问题折磨按本文的思路先做一轮小规模验证大概率会发现之前不敢想的上下文长度其实离你只差一个 4 bit 的距离。