做长上下文推理优化的时候估计不少人跟我一样被KV Cache这块硬骨头卡过上下文一拉长显存先爆接着访存带宽拖后腿decode阶段慢得让人抓狂。我最近在一个内部代号为“202606”的项目里试了试MXFP4方案把注意力计算里的KV缓存改成了MX格式的4-bit浮点再配合自定义的融合算子效果比预想的要稳。这里把我踩过的坑、调参的心得和实测数据整理出来给同样在搞大模型推理、量化部署或者算子开发的朋友一个参考。先说清楚这文章聊的是什么MXAttention是我在项目里做的一个面向Transformer解码器的融合注意力实现核心低精度格式是MXFP4也就是OCP MX规范里的4-bit浮点格式。它解决的是三个问题KVCache的显存占用、解码阶段的内存带宽瓶颈、以及低精度算子在注意力这个场景下的精度退化。适合规模部署大模型、在长上下文下做推理优化、或者研究低精度训练/推理的工程师读。1. 先理解MXFP4到底是什么它跟INT4有什么区别1.1 FP4不是INT4的替代品它走的是另一条路很多人一看到4-bit精度第一反应是“这不就是INT4吗”。我一开始也是这么想的真正用下去才发现两者思路完全不同。INT4是固定小数点格式把数值范围均分成16个阶梯均匀量化。FP4是浮点格式走的是科学计数法的路子1位符号、若干位指数、若干位尾数。典型配置有两种e2m1和e3m0。e2m1意味着2位指数加1位尾数能表达的最大值到6最小的正常浮点数在0.5附近非正常值可以下探到0.125左右动态范围跨度远大于INT4。e3m0则是3位指数加0位尾数相当于只有8个不同的量级阶没有尾数精度只能表示1、2、4、0.5、0.25这些二的幂次。这个差异看起来只是格式问题放在量化场景里就是本质区别。注意力里Q、K、V向量经过LayerNorm之后分布虽然相对稳定但随着位置编码叠加和跨层传播数值范围还是会波动的尤其在长上下文下某些token位置的K/V范数可能比平均值大好几倍。INT4用均匀阶梯去切这个分布遇到极值就容易产生严重的相对误差。FP4的指数编码天然提供了更大的动态范围即使尾数只有1位也能通过指数偏移把量程拉开。我实际测过一组数据在一个7B模型上同一批KV向量INT4量化带来的attention score误差在部分head上能到4%以上MXFP4的e2m1配置普遍压在1.5%以内。差异主要来自对异常值的刻画能力。如果向量内部元素都比较接近INT4精度更好因为尾数位宽更足但只要分布出现拖尾FP4的优势就出来了。注意力里的KV向量恰恰经常有拖尾这就注定了MXFP4在这个场景比INT4更适用。1.2 MX共享Scale的设计逻辑MXFP4通常不会单独使用而是配合MX规范里的block scale机制一起出现。所谓block scale就是每固定数量的元素共享一个缩放因子论文和硬件实现里常见的是32个元素作为一个block每个block配一个FP32的scale。实际数值等于FP4尾数乘以对应的scale。这个设计的本质是为了解决浮点格式在数值太大或太小时的表达空白。如果你不用block scale直接拿e2m1去量化任意向量假设某个元素是300而FP4最大只能表达6那就只能截断成6信息直接丢失。但如果每32个元素配一个scale比如把scale设成50那300除以50等于6正好落在FP4的表达范围内反量化时乘回50300就保住了。FP4负责表达相对大小和精细差值scale负责适配量级两者结合就让4-bit格式拥有了接近动态范围无限的表达能力。MX规范里对scale的选取还有讲究通常用的是矩阵或向量的绝对最大值或者二阶范数。我在实现里用的是amax机制先取block内32个元素的绝对值最大值再按FP4可表示的最大浮点值去反推scale。有一点容易漏因为e2m1的最大值是6所以实际scale要取max_value/6而不是直接取max_value。如果直接用max_value当scale编码时所有值除以scale之后落在0到1之间白白浪费了FP4在1到6之间的全部表达空间精度亏得很明显。我最初就是在这里没注意导致Kernel跑起来困惑度一下掉了0.3查了半天才发现是scale计算公式少除了一个6。1.3 什么时候选MXFP4而不是FP8或者INT8用MXFP4之前我其实先试过标准的FP8 KV Cache。FP8的精度确实不错训练和推理都比较友好但KV Cache的显存开销只减了一半与4-bit方案相比没有质的飞跃。在GQAGrouped Query Attention结构下KV Cache的体量本来就比MHA小一些但如果上下文长度上到32K、64KFP8依然会吃掉大量HBMdecode带宽问题也依然存在。另一种思路是INT8加per-channel scale或者做比较激进的INT4加per-group scale。INT8方案误差控制好但显存压缩率不够INT4加group scale的方案能压下去可对一个group内的异常值鲁棒性差。我的经验是如果模型本身的激活分布方差比较大或者目标上下文超过16KMXFP4在精度和压缩率的平衡上是更优解如果模型输出非常稳定或者主要跑短上下文FP8或INT8反而实现成本更低生态工具也更成熟。2. 为什么注意力计算盯上了KV Cache2.1 真正贵的是访存不是计算看Attention的FLOPs很多人第一反应是“这玩意儿算力需求很大”但从实际Kernel profiling来看decode阶段的瓶颈几乎都在内存带宽上。单token生成时Q只有一个向量K和V是整个序列的缓存规模是序列长度乘以隐藏维度再乘以层数。计算量很小但K/V的读取量是逐token线性增长的。假设序列长度是4096隐藏维度是4096单层KV缓存就是4096乘以4096FP16下要32MB模型有32层就是1GB。这个数据要从HBM读进SM比例远超计算量本身。把KV压成MXFP4之后同样的数据量变成原来的四分之一HBM读压力直接下降。GPU上很多算子的效率卡在DRAM带宽上只要压了内存搬运量带来的加速比是实打实的。我在A100上实测decode阶段把KV从FP16换成MXFP4在其他条件不变的情况下单token延迟下降了大约35%到40%换算成端到端吞吐提升非常可观。还有一个容易被忽略的点L2 Cache。FP16的KV数据占L2空间太大导致下一层要读的数据经常被挤出L2反复miss回HBM。压成4-bit后同样大小的L2能装下四倍序列长度的KVcache命中率明显改善。尤其在中短序列下这种L2层面的改善比HBM带宽的改善更明显。2.2 融合算子要从数据流角度重新拆解MXAttention不是我简单地把KV cache换了个存储格式然后把反量化逻辑塞进现有的FlashAttention里而是从数据流层面重新拆了一遍Attention实现。标准Attention流程是读Q和K算QK^TSoftmax缩放再乘V得到输出。FlashAttention把这个过程做了分块和在线Softmax优化减少了中间矩阵的落盘。MXAttention在这个基础上又加了一层KV的读取和反量化与后续矩阵乘融合让FP4数据直接进寄存器或共享内存。具体来说我按block加载FP4格式的K先把scale和尾数都读进寄存器在寄存器里做反量化还原出FP16或者FP32的值立刻参与QK^T计算。这样FP4数据只在HBM里占4-bit空间一旦进入芯片内部就在寄存器层面恢复成高精度浮点避免了把整个KV Cache先反量化到一个临时FP16缓冲区再计算的额外访存。这个流程看着没什么但在长序列下差别很大如果不做融合反量化后写回临时缓冲区就得多一次HBM写入和一次HBM读取一来一回等于多搬了8倍数据。2.3 从Prefill到Decode的不同策略Prefill阶段和Decode阶段对MXFP4的收益差异也很大。Prefill阶段并行处理一整段promptQ矩阵很大计算密度高此时KV Cache访存占比相对低单纯的访存优势不明显反而是精度损失带来的质量影响更容易暴露。因此我在实际部署时Prefill阶段默认走FP16 KV只有Decoder阶段才启用MXFP4。两个阶段共用同一份KV存储会有点麻烦我的做法是写了个kernel入口根据phase参数决定是否启用MXFP4 loadPrefill写FP16Decode读的时候如果发现存储格式是MXFP4就走反量化路径如果有必要甚至可以在Prefill完成之后把FP16的KV原地转成MXFP4。Decode阶段为什么值得用MXFP4因为此时每一步只有一个token的QK/V的读取成为绝对主导访存压缩的收益最大。实际跑起来同一个模型同一个prompt混合策略的pipeline在吞吐上比全程FP16高出约1.8倍而困惑度几乎持平相差不到0.05。3. 核心实现细节与量化方案3.1 按Block量化KV的在线Scale计算MXFP4的量化并不是提前离线算好的而是在推理过程中实时计算。每次KV Cache需要写入时计算它的block scale然后把尾数写进缓存。这个流程适合叠加到原始的KV Cache写入kernel里不额外增加一次全量扫描。具体步骤我整理成了这个流程取当前KV向量的一段长度32按block切分。计算该block的amax即绝对值最大元素。根据amax和FP4最大可表示值6算出scale amax / 6。向量里的每个元素除以scale就近取整到FP4可表示的浮点集合。把scale用FP32保存下来与4-bit尾数一起存储。这里有几个容易踩的细节。一是只要遇到amax为0的blockscale直接置为0反量化时把整个block视为0不用走除法性能和安全都兼顾。二是在计算scale时如果直接用amax而不是amax除以6会损失尾数精度前面说过了。三是block切分最好按最后一个维度连续切这样才能在硬件上连续读取避免cross-stride的scatter操作。从误差角度看per-block scale比per-tensor scale强得多。per-tensor scale相当于让整条KV共享一个缩放因子为了让极端值不溢出scale会偏大大部分普通元素编码后都挤在很小的量级上精度损失明显。per-block scale只有32个元素共享一个scale局部动态范围适配得好这也是MX规范推荐这种粒度的原因。3.2 QK^T、Softmax和Aggregation的全流程量化不是只换个存储格式就完了关键是计算过程怎么组织。我实现的MXAttention流程是这样的Q保持FP16或者FP32不量化。K和V在HBM里以MXFP4格式存储。计算QK^T时按block读取K反量化到FP16或FP32然后和Q做点积。Softmax的结果在线计算我用的是FlashAttention风格的online softmax维护running max和running sum。之后O softmax(QK^T)VV同样从MXFP4格式按block读取并反量化。累加过程全程用FP32。这个流程里最需要注意的点是Q不要压成FP4。我试过把Q也压到MXFP4推理质量下滑得厉害尤其是一些对数值敏感的head几乎直接失效。原因是Q和K做点积时Q的精度直接影响每个注意力权重的精度一旦Q量化出现相对误差相当于给score加了个噪声。把Q保持在FP16K用FP4相当于只对内存体积大的那一侧做压缩而计算精度主要依赖Q。这是性价比非常高的折中方案。3.3 关键Kernel伪代码示例我贴一个简化版的kernel逻辑便于理解整体结构。这个版本忽略了很多边界条件和硬件细节但数据流是对的。实际开发时我是先用这个简单版本跑通正确性再去写TMA和双缓冲优化版本。// 简化版MXFP4 KV Cache Attention // Q: [num_heads, head_dim] FP16 // K_mx: [seq_len, head_dim/2] MXFP4 (因为head_dim一般128, block_size32) // K_scale: [seq_len, head_dim/32] FP32 __global__ void mx_attention_decode_kernel( const half2* Q, const uint8_t* K_mx, // 4bit packed const float* K_scale, const uint8_t* V_mx, const float* V_scale, float* O, int seq_len, int head_dim) { int tid threadIdx.x; int block_seq blockIdx.x; // 按序列分块 float q_local[HEAD_DIM]; // load full Q to registers (FP16 - FP32) load_half2_to_float(Q, q_local, HEAD_DIM); float acc[HEAD_DIM] {0.0f}; float m_i -1e30f; float l_i 0.0f; for (int kb 0; kb seq_len; kb BLOCK_SIZE) { // 读取K的MXFP4 block uint8_t k_packed[32 / 2]; float k_scale K_scale[kb / 32]; load_mxfp4_block(K_mx kb * head_dim / 2, k_packed, BLOCK_SIZE); // 反量化K float k_vec[HEAD_DIM]; dequant_mxfp4_block(k_packed, k_scale, k_vec, HEAD_DIM); // 计算QK^T 点积 float score dot_product(q_local, k_vec, HEAD_DIM); // online softmax float m_new fmaxf(m_i, score); float alpha __expf(m_i - m_new); float p __expf(score - m_new); l_i l_i * alpha p; // 读取V block 并累积 uint8_t v_packed[32 / 2]; float v_scale V_scale[kb / 32]; load_mxfp4_block(V_mx kb * head_dim / 2, v_packed, BLOCK_SIZE); float v_vec[HEAD_DIM]; dequant_mxfp4_block(v_packed, v_scale, v_vec, HEAD_DIM); for (int d 0; d HEAD_DIM; d) { acc[d] acc[d] * alpha p * v_vec[d]; } m_i m_new; } // 归一化 for (int d 0; d HEAD_DIM; d) { O[d] acc[d] / l_i; } }实际工程版本里有两处变化比较大一是head_dim通常是128一个矩阵块同时覆盖多个blockscale就不是单个值了而是一个小数组二是TMA可以一次性把FP4数据按block拉进SMEM避免逐元素读取。代码里的dequant函数就是按block scale对32个元素逐一乘回去。反量化本身开销很低就是一次乘法加一次格式转换关键是别把这步放到HBM里做。3.4 显存布局和Packing细节MXFP4有一个工程上麻烦的地方它不是字节对齐的格式。一个元素4-bit一个字节能塞两个元素。如果直接按字节存取必须考虑packing顺序。我用的是低位在前即第一个元素存在字节的低4位第二个元素存在高4位。这个约定需要在写KV和读KV两边保持一致否则一路错到底。另外block size选32不只是为了精度也是为了对齐。32个FP4元素正好占16个字节等于一个float4的尺寸读取时可以做128-bit memory transaction带宽利用最充分。如果block size改成16字节数变成8读取效率差一些如果选64scale粒度变粗精度下降。我对比过32和64的偏差64的困惑度比32平均差0.08左右但访存效率反而没明显提升所以最终固定为32。4. 实测数据与参数选型4.1 与不同KV缓存方案的精度对比我在一个内部微调过的7B模型上做了对比实验评测时用了困惑度指标和一组人工设计的极端prompt集。同一模型、同一输入KV Cache分别用FP16、FP8-E4M3、INT4-Group和MXFP4-e2m1存储。KV Cache格式每元素比特数相对FP16显存困惑度越低越好Attention Score最大偏差FP16161x8.720%FP8 E4M380.5x8.760.3%INT4 Group3240.25x9.052.8%MXFP4 e2m140.25x8.811.4%FP16的困惑度最低这是预期中的。FP8的精度损失极小但如果只看显存收益它没有4-bit方案的吸引力。INT4的困惑度恶化最明显而且位置靠后的层误差有累积效应。MXFP4的困惑度比FP8高了0.05左右比INT4低了0.24同时拿到4倍压缩综合看是最划算的。我后来又测了e3m0配置困惑度掉到了9.3直接放弃e3m0虽然表示范围更大但没有尾数位对普通数值的刻画太粗糙了。4.2 吞吐和显存收益集成MXAttention后我在同样的推理框架下对比了端到端的效果。模型配置是7B、GQA、32层输入prompt 4096 token输出512 token测了FP16 baseline和MXFP4两种情况。场景KV显存占用Decode吞吐token/s单层Kernel耗时占比FP16 baseline12.6GB64.5100%MXFP4 KV3.2GB89.272%MXFP4 融合3.2GB104.858%这里有三个结论值得细看。第一显存从12.6GB降到3.2GB意味着同样的显存预算下上下文长度可以拉长到原来的约4倍。第二光换格式不融合decode吞吐提升大约38%这部分提升主要来自HBM带宽压力的下降。第三加上融合kernel后在相同存储格式下又涨了将近18%这部分收益来自避免反量化中间缓冲区的额外写入和读取。第三个结论可能出乎不少人的预期但实际上长序列下中间缓冲区的读写量非常大省掉这一步带来的加速比一点不比压缩带宽少。4.3 参数选型建议根据我实验的积累MXFP4用于注意力的参数选型顺序应该是block_size优先设32格式优先选e2m1Q不量化scale用FP32保存。这四个参数基本不用再纠结。block_size如果设成16scale多了精度微升但读取效率下降实测吞吐损失约7%。设成64精度下降明显但吞吐提升只有2%左右不值。e2m1和e3m0之间不用犹豫除非你对显存有极端要求且能接受较大精度损失否则e2m1是唯一合理的选择。scale的存储精度我试过FP16和FP32FP16 scale会让某些大头位置出现0.5%左右的额外误差而FP32 scale占用只多了一点点因为scale数量是元素数的三十二分之一成本可忽略所以无脑选FP32。另外还有一个容易被忽略的点某些模型做了Sliced Attention或者GQA分组之后KV在不同group之间重复存储让量化收益打了折扣。如果遇到这种结构建议在KV写入层就把分组的共享关系利用起来相同数据只存一份MXFP4。5. 常见问题与排查技巧实录5.1 数据流中莫名其妙的NaN和Inf我第一次把MXAttention跑起来的时候前向直接出现NaN排查了整整一个下午。后来发现根因不在kernel本身而是scale计算的一个边界条件某些padding位置的向量全为0amax是00除以6等于0scale是0反量化时x乘以0等于0这本身没问题。但问题是有些硬件路径在scale为0的block里如果恰好有非零尾数做除法时会触发invalid operation直接产生NaN。解决方案有两层。第一层是在量化阶段遇到amax是0的block直接写一个特殊的scale值我用的-1.0做标记反量化时看到-1.0就直接返回0不走乘法路径。第二层是在kernel里加clamp所有反量化后的值都限制在一个合理范围内。这两层让我在后面几个模型上再也没见过NaN。5.2 收益没有想象中大反量化开销补偿问题有朋友抄了我的方案之后反馈说速度没提升甚至变慢了。他用的实现是先开一个FP16缓冲区把MXFP4的KV全部反量化到缓冲区然后再跑标准Attention。这个做法等于你省了HBM读4-bit但又把反量化后的FP16写回HBM再读一遍。写入165MB、再读165MB加在一起比直接读FP16还要多出一倍流量自然慢。正确的做法必须是反量化和计算融合在一个kernel里完成FP4数据从HBM读出来后直接留在寄存器或SMEM里经过scale恢复成高精度值就立刻参与计算绝无第二次落地的过程。这个道理说起来简单但实际写TMA版本的时候很容易为了图省事先搞一个临时缓冲区。我自己的经验是可以先用临时缓冲区版本验证正确性和精度再花时间做真正融合的版本因为后面这个版本的收益占到总收益的三四成。5.3 长上下文尾部的精度退化在长上下文场景下我观察到decode到后半段时模型输出质量比起前半段略有下降困惑度升高了约0.15。排查后发现不是MXFP4的问题而是attention得分在尾部极端情况下出现了较大的绝对值。当某个query和某个key的内积远大于正常分布时softmax会把概率几乎全给它此时这个score即使只有1%的误差也会对最终输出产生明显影响。解决方案是给MXFP4的K增加一个per-block的截断策略。当block的amax超过某个阈值时依然用FP4存储但在计算QK^T时对score做一个logit cap限制极端score的继续增大。这个操作不影响绝大多数token的正常计算只在极端情况下保持数值稳定。我在这里多花了大概200行代码换来了长上下文下稳定不掉点值得。5.4 不同算子库之间的格式兼容MXFP4目前还不是所有深度学习框架都能原生支持的类型。PyTorch里没有内置的MXFP4 dtypeTensorRT的某些版本对MX格式支持也仅限于部分算子。在工程落地时我建议把MXFP4当作一种“私有的紧凑存储格式”来看待而不是当作框架原生类型。KV Cache写入时由自定义kernel负责编码读的时候由自定义kernel负责解码框架本身只要保证能搬运uint8的tensor即可。如果你要跟别的组件共享KVCache比如投机解码或者多轮对话的缓存管理需要特别注意格式标记。我在工程里用一个额外的int32 tensor记录了每个block的scale状态并在KV缓存管理结构里增加一个format字段值为0表示FP16、1表示MXFP4。在切换格式时需要重新初始化KV Cache。这个字段虽然小但避免了不同阶段误读数据导致的诡异错误。6. 一些从实测打磨出来的工程经验如果让我只挑几条最重要的经验说我会挑这三条。第一MXFP4方案不要试图在所有阶段统一启用Prefill和Decode要区别对待Prefill阶段老老实实用FP16Decode阶段再用MXFP4既保住质量又不牺牲速度。第二block scale必须和反量化kernel严格绑定中间任何一次格式转换都可能让scale语义失效KV Cache的生命周期内不要动格式。第三先做一个能跑通的低性能版本把精度验证了再去做TMA等高性能优化不要一上来就追求极致效率否则出问题时很难定位是格式问题还是kernel问题。我在调MXAttention的过程中还有一个体会注意力算子看似简单实际上每个细节都在跟访存和数据布局较劲。MXFP4的4倍压缩是敲门砖真正拉开差距的是把反量化、load、计算融合成一个完整的数据流水。这个思路不但适用于KV Cache也适用于其他任何大模型推理中的内存密集算子。项目做到后面我已经在考虑把同样的MXFP4策略用在FFN的中间激活上不过那是另一个话题了。