Research map / Architectures and efficiency / Attention and transformers

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

NeurIPS 2026 · Architectures and efficiency › Attention and transformers · approach: Sparse attention for long-context LLMs · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes SnapMLA, an FP8 decoding framework for DeepSeek-style MLA that keeps the RoPE part in high precision, uses per-token KV quantization, and reconstructs the FP8 PV pipeline with scale fusion, reporting up to 1.91x decoding throughput with near-BF16 benchmark quality.

Problem

FP8 attention for the decoding phase of MLA-based LLMs is hard because the decoupled RoPE part is numerically sensitive, per-block KV scales conflict with autoregressive decoding, and per-token V scales are misaligned with the FP8 PV GEMM k-major layout.

Approach

Keep the MLA RoPE component in BF16 while quantizing the latent content cache per token; pre-scale RoPE queries and keys into the FP8 quantization domain so the QK kernel accumulates uniformly; fuse the per-token V scale into the probability block P and fold block-wise P scales into online softmax accumulation so PV needs no post-GEMM dequantization.

Key results

• Up to 1.91x end-to-end decoding throughput over BF16 FlashMLA on long-output decoding workloads (16k to 128k context, DP/TP configurations).
• FP8 SnapMLA is near parity with BF16 FlashMLA on DeepSeek-V3.1, e.g. MMLU-Pro 84.43 vs 84.41 and AIME-24 93.65 vs 93.85.
• On LongCat-Flash-Thinking, FP8 scores are close to BF16, e.g. LCB 79.74 vs 79.40 and IFEval 87.80 vs 86.90, with some small drops such as AIME-25 88.44 vs 89.58.

Surprising findings

• Including the RoPE part in per-token quantization causes a large error increase in deeper layers, supporting the RoPE-aware design.
• Some FP8 scores slightly exceed BF16 on certain benchmarks (e.g., MMLU-Redux 90.89 vs 90.48 on DeepSeek-V3.1).

Limitations

• Kernel speedups are specific to NVIDIA Hopper GPUs and would need adaptation for other hardware.
• Evaluation covers only decoding and quantizes only E4M3; longer contexts and more aggressive formats are unexplored.
• Hopper GPU model is not disclosed, limiting reproducibility of absolute numbers.

Experimental setup

Datasets: MMLU-Pro MMLU-Redux IFEval Arena-Hard MATH-500 HMMT-25 AIME-24 AIME-25 BeyondAIME GPQA-Diamond ZebraLogic LiveCodeBench

Benchmarks: MMLU-Pro MMLU-Redux IFEval Arena-Hard MATH-500 AIME-24 AIME-25 BeyondAIME GPQA-Diamond ZebraLogic LiveCodeBench

Metrics: accuracy Mean@k generation throughput TFLOPS RMSE cosine difference relative L2 error

Baselines: FlashMLA BF16

Builds on: FlashAttention-3 FlashMLA DeepSeek-V2 MLA

Closest papers by meaning

• SnapKV: LLM Knows What You are Looking for Before Generation (NeurIPS 2024 · similarity 0.88)
• Progressive Mixed-Precision Decoding for Efficient LLM Inference (ICLR 2025 · similarity 0.87)
• QuantSpec: Self-Speculative Decoding with Hierarchical Quantized KV Cache (ICML 2025 · similarity 0.86)
• KVTuner: Sensitivity-Aware Layer-Wise Mixed-Precision KV Cache Quantization for Efficient and Nearly Lossless LLM Inference (ICML 2025 · similarity 0.86)
• CommVQ: Commutative Vector Quantization for KV Cache Compression (ICML 2025 · similarity 0.86)
• Quamba: A Post-Training Quantization Recipe for Selective State Space Models (ICLR 2025 · similarity 0.86)