arXiv:2602.10718cs.LGcs.CL2026-02被引 2

通过硬件感知的FP8量化流水线,显著提升长文本生成效率。

SnapMLA: Efficient Long-Context MLA Decoding via Hardware-Aware FP8 Quantized Pipelining

  • 针对MLA架构设计分词精度保留的量化策略
  • 实现1.91倍长输出解码吞吐量提升,精度几乎无损
  • 适合追求高吞吐长文本生成的模型部署场景

尽管FP8注意力在FlashAttention-3等创新中展现出巨大潜力,但将其集成到DeepSeek多头潜在注意力(MLA)架构的解码阶段仍面临挑战,包括位置编码解耦引发的数值异质性、FP8 PV GEMM中量化尺度错位,以及系统级优化支持需求。本文提出SnapMLA,一种面向长上下文高效解码的FP8 MLA框架,通过以下硬件感知的算法-内核协同优化:(i) RoPE感知的逐令牌KV量化:基于对MLA KV缓存异质量化敏感性的分析,该方法保留RoPE部分的高精度,并采用逐令牌粒度以匹配自回归解码过程,维持量化准确性;(ii) 量化PV计算流水线重构:解决因MLA共享KV结构导致的FP8 PV计算中量化尺度错位问题;(iii) 端到端数据流优化:采用专用内核建立高效读写流程,保障数据流动顺畅,提升性能。在主流MLA大模型上的大量实验表明,SnapMLA在长输出解码工作负载下,吞吐量最高提升1.91倍,且在推理与代码生成基准测试中,与BF16基线相比保持近似一致的基准质量。代码已开源:https://github.com/meituan-longcat/SGLang-FluentLLM。

原文摘要 · Abstract (English)

While FP8 attention has shown substantial promise in innovations like FlashAttention-3, its integration into the decoding phase of the DeepSeek Multi-head Latent Attention (MLA) architecture presents notable challenges. These challenges include numerical heterogeneity arising from the decoupling of positional embeddings, misalignment of quantization scales in FP8 PV GEMM, and the need for optimized system-level support. In this paper, we introduce SnapMLA, an FP8 MLA decoding framework optimized to improve long-context efficiency through the following hardware-aware algorithm-kernel co-optimization techniques: (i) RoPE-Aware Per-Token KV Quantization: Motivated by our analysis of the heterogeneous quantization sensitivity inherent to the MLA KV cache, this approach preserves the RoPE part in high precision. Furthermore, per-token granularity is employed to align with the autoregressive decoding process and maintain quantization accuracy. (ii) Quantized PV Computation Pipeline Reconstruction: Addresses the misalignment of quantization scales in FP8 PV computation caused by the shared KV structure of the MLA. (iii) End-to-End Dataflow Optimization: Establishes an efficient data read-and-write workflow using specialized kernels, ensuring streamlined data flow and improved performance. Extensive experiments on state-of-the-art MLA LLMs show that SnapMLA achieves up to a 1.91x improvement in throughput on long-output decoding workloads while maintaining near-parity benchmark quality compared with the BF16 baseline on the evaluated reasoning and code-generation benchmarks. Code is available at https://github.com/meituan-longcat/SGLang-FluentLLM.

FP8量化长文本生成MLA架构硬件优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。