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, where the RoPE part is maintained in high precision, motivated by our comprehensive analysis of the heterogeneous quantization sensitivity inherent to the MLA KV cache. Furthermore, per-token granularity is employed to align with the autoregressive decoding process and maintain quantization accuracy. (ii) Quantized PV Computation Pipeline Reconstruction, which resolves the misalignment of quantization scale in FP8 PV computation stemming from the shared KV structure of the MLA KV cache. (iii) End-to-End Dataflow Optimization, where we establish an efficient data read-and-write workflow using specialized kernels, ensuring efficient data flow and performance gains. Extensive experiments on state-of-the-art MLA LLMs show that SnapMLA achieves up to a 1.91x improvement in throughput, with negligible risk of performance degradation in challenging long-context tasks, including mathematical reasoning and code generation benchmarks. Code is available at https://github.com/meituan-longcat/SGLang-FluentLLM.

翻译：尽管FP8注意力机制已在FlashAttention-3等创新中展现出巨大潜力，但其在DeepSeek多头潜在注意力（MLA）架构解码阶段的集成仍面临显著挑战。这些挑战包括：位置嵌入解耦导致的数值异质性、FP8 PV GEMM中量化尺度失准，以及对优化系统级支持的需求。本文提出SnapMLA——一种经优化的FP8 MLA解码框架，通过以下硬件感知的算法-内核协同优化技术提升长上下文处理效率：（i）RoPE感知的逐令牌KV量化：基于对MLA KV缓存固有异质量化敏感性的全面分析，将RoPE部分保持在高精度状态；同时采用逐令牌粒度以匹配自回归解码过程并维持量化精度。（ii）量化PV计算流水线重构：解决因MLA KV缓存的共享KV结构导致的FP8 PV计算中量化尺度失准问题。（iii）端到端数据流优化：通过专用内核建立高效数据读写工作流，确保数据流动效率与性能增益。在先进MLA大语言模型上的大量实验表明，SnapMLA在具有挑战性的长上下文任务（包括数学推理与代码生成基准测试）中，吞吐量最高可提升1.91倍，且性能退化风险可忽略。代码发布于https://github.com/meituan-longcat/SGLang-FluentLLM。