---
title: 'SnapMLA: FP8 Decoding for MLA Models'
url: https://www.emergentmind.com/topics/snapmla
type: topic
---

# SnapMLA: FP8 Decoding for MLA Models

SnapMLA is an FP8 decoding framework for long-context inference in DeepSeek-style Multi-head Latent Attention (MLA) language models. It targets the specific numerical and systems problems that arise when MLA is executed in low precision, especially under absorbed inference mode, where content K/V are represented through a shared latent cache while RoPE is decoupled. Its design combines three hardware-aware algorithm–kernel co-optimization techniques—RoPE-Aware Per-Token KV Quantization, Quantized PV Computation Pipeline Reconstruction, and End-to-End Dataflow Optimization—and is reported to achieve up to a \(1.91\times\) improvement in throughput, with negligible risk of performance degradation in challenging long-context tasks including mathematical reasoning and code generation benchmarks [2602.10718].

## 1. Terminology and research context

In the SnapMLA literature, **MLA** denotes **Multi-head Latent Attention**, the attention architecture used in DeepSeek-V2/V3 and the LongCat-Flash series. In this setting, MLA compresses the KV cache by storing a low-rank latent representation rather than per-head, full-dimensional K/V tensors, and SnapMLA concerns the **decoding** path for such models [2602.10718].

The acronym **MLA** is not unique across machine learning. A separate multimodal-learning line uses MLA to mean **Multimodal Learning with Alternating Unimodal Adaptation**. That work studies supervised multimodal classification with modality-specific encoders, a shared head, gradient modification, and uncertainty-based fusion; its source text explicitly states that “SnapMLA” is not explicitly mentioned there, and that such a term would plausibly refer to a practical implementation of that multimodal MLA formulation [2311.10707]. This terminology overlap is substantial enough that disambiguation is often necessary in technical discussion.

A second relevant development is **TransMLA**, which formalizes conversion of GQA-based pretrained transformers into MLA-based models. TransMLA defines MLA as a low-rank latent KV mechanism with per-head up-projection, proves that MLA is more expressive than GQA when both have the same size of KV cache, and positions such conversion as the architectural basis for a “SnapMLA”-type system on mainstream LLaMA-, Qwen-, or Mistral-style checkpoints [2502.07864]. This suggests a practical lineage in which TransMLA broadens the set of models that can host SnapMLA-style decoding kernels, while SnapMLA itself focuses on efficient long-context execution for already-MLA-structured models.

## 2. Multi-head latent attention as the computational substrate

MLA, in the DeepSeek/LongCat sense, stores a shared latent vector for K and V rather than materializing full per-head K/V states for every token. For a token hidden state \(h_t \in \mathbb{R}^d\), the latent cache component is

\[
c_{KV} = W^{DKV} h_t,\quad c_{KV} \in \mathbb{R}^{d_c},\ d_c \ll d.
\]

From this latent, the model derives content key and value components, while the positional component is handled separately through RoPE:

\[
k_i^C = W^{UK} c_{KV},\qquad
k^R = \mathrm{RoPE}(W^{KR} h_t),\qquad
k_i = [k_i^C; k^R],\qquad
v_i^C = W^{UV} c_{KV}.
\]

The resulting KV cache stores, for each token, the low-rank latent vector \(c_{KV}\), shared between all heads and between K/V, together with a RoPE key component \(k^R\), which is shared across heads but not compressed. This organization is the basis for MLA’s lower cache footprint relative to standard MHA [2602.10718].

SnapMLA is formulated around the **inference-optimized absorbed mode** of MLA, in which K/V up-projections are folded into query and output projections so that \(k_i^C\) and \(v_i^C\) do not need to be explicitly reconstructed during decoding. In this mode, the attention score for query head \(i\) at time \(t\) against a past token \(j\) is

\[
q_{t,i}^\top k_{j,i}
=
\underbrace{(q_{t,i}^C W^{UK})^\top c_{KV,j}}_{\text{content term}}
+
\underbrace{(q_{t,i}^R)^\top k_j^R}_{\text{RoPE term}}.
\tag{1}
\]

This decomposition is computationally attractive because the dominant historical state is the shared latent cache, but it also creates a heterogeneous numerical structure: the content term is mediated by the latent \(c_{KV}\), whereas the RoPE term is carried by a decoupled positional vector. SnapMLA is specifically designed for this absorbed decoding regime rather than for generic attention kernels [2602.10718].

## 3. Failure modes of naïve FP8 MLA decoding

The central problem SnapMLA addresses is that MLA’s structural compression does not by itself guarantee efficient **FP8 decoding**, particularly at long context lengths. The paper identifies three obstacles: numerical heterogeneity from decoupled RoPE, quantization-scale misalignment in FP8 PV GEMM, and system-level support requirements [2602.10718].

The first obstacle is **heterogeneous quantization sensitivity**. In the reported analysis for LongCat-Flash-Thinking, content components derived from \(c_{KV}\) are tightly concentrated around zero, approximately within \(\pm 10^1\), whereas RoPE components \(k^R\) have a much wider dynamic range, up to \(\pm 10^3\), with heavy tails and outliers. Under uniform FP8 quantization, mean-squared error on the RoPE part increases sharply, and the resulting distortion accumulates across layers, especially deeper ones. Because RoPE encodes positional structure, this degradation is especially harmful for long-context behavior.

The second obstacle is **scale misalignment in PV computation**. In MLA, V is derived from the same latent cache as K and therefore inherits **per-token scales** from \(c_{KV}\). On Hopper, however, FP8 PV GEMM requires a **k-major layout** for V, meaning the value tensor must be contiguous along the sequence dimension that also acts as the reduction dimension. This couples the per-token quantization scales to the reduction axis in a way that is incompatible with standard post-GEMM dequantization strategies. A naïve FP8 implementation either loses accuracy or incurs expensive rescaling overhead.

The third obstacle is the overall **latency-sensitive, memory-bound** nature of autoregressive decoding. Even if FP8 reduces storage cost, its benefit can be offset by transpositions, separate quantize/dequantize passes, fragmented memory layouts, or synchronization-heavy mixed-precision kernels. SnapMLA treats these issues as joint algorithmic and kernel-level design constraints rather than as isolated implementation details [2602.10718].

## 4. RoPE-aware quantization and scale-fused attention

SnapMLA’s first major design decision is **RoPE-aware per-token KV quantization**. It quantizes only the **content** portion of the KV state, \(c_{KV}\), and does so at **per-token granularity**, while keeping the RoPE component \(k^R\) in **BF16**. The paper motivates per-token quantization on two grounds: it aligns with the autoregressive token-by-token append process, and it better captures token-level outliers than coarser per-tensor or per-block schemes [2602.10718].

Per-token quantization is introduced in the standard form

\[
X_q^{(i,:)} = \mathrm{round}\Big(\frac{X^{(i,:)}}{S_i}\Big),
\]

with one scale \(S_i\) for each token row. The RoPE part is excluded from FP8 quantization because quantizing RoPE directly was shown to introduce large RMSE, cosine error, and relative \(L_2\) error, with layer-wise “error explosion” in deep layers.

Keeping RoPE in BF16 creates an additional mismatch, because QK GEMM would then need to combine FP8 content terms and BF16 RoPE terms in the same reduction. SnapMLA resolves this by **pre-aligning the numeric domains** through RoPE pre-scaling:

\[
Q^R \leftarrow \frac{Q^R}{S^{Q_c}},\qquad
K^R \leftarrow \frac{K^R}{S^{K_c}}.
\]

This allows the GEMM pipeline to treat the RoPE contribution as numerically aligned with the FP8 content contribution without introducing a mixed-precision accumulation schedule that would disrupt the original kernel pipeline.

The second major design decision is **Quantized PV Computation Pipeline Reconstruction**. If V is represented as

\[
V = V_q \cdot S_V,
\]

then the output

\[
O = PV
\]

can be rewritten by fusing the per-token value scale into the probability matrix:

\[
P' = P \cdot S_V,\qquad
O = P' V_q.
\]

This transfers the scale handling from V to P, making the FP8 PV GEMM compatible with Hopper’s k-major constraints. Because this scale fusion enlarges the dynamic range of \(P'\), SnapMLA applies **block-wise dynamic quantization** to \(P'\), with block size \(BlockN = 64\), and integrates scale handling into an **online softmax** update rather than relying on a global post-GEMM dequantization. The paper describes this as implicit dequantization through the normalization logic itself.

A further issue appears in the dual-warp-group pipeline inherited from FlashMLA-style scheduling: processing blocks out of order can induce **scale hazards**, because accumulated outputs become expressed in the scale of a later block before earlier contributions are incorporated. SnapMLA reconstructs the scheduling to enforce a monotonic order for the relevant warp group, avoiding unstable back-and-forth rescaling [2602.10718].

## 5. Hardware-aware kernel pipeline

SnapMLA is explicitly designed for **NVIDIA Hopper**. It uses **FP8 E4M3** for activations and KV cache content, **FP8 Tensor Cores (WGMMA)**, **TMA (Tensor Memory Accelerator)**, **Swizzle-128B shared-memory layouts**, asynchronous TMA-to-SMEM staging, and warp-specialized pipelines that separate compute and data movement [2602.10718]. The framework is not simply an FP8 variant of an existing MLA kernel; it is a restructuring of the decoding dataflow around Hopper’s memory hierarchy and WGMMA layout rules.

At the **model level**, SnapMLA leaves the MLA architecture unchanged. At the **algorithmic level**, it introduces the RoPE-aware quantization and scale-fused attention mechanisms described above. At the **kernel level**, it implements a set of fused operators. The **Fused Q-Quant and K-Append** kernel performs per-token scale calculation for Q and K content, FP8 quantization of content, mixed-precision conversion for RoPE, RoPE scale-domain alignment, and direct append of quantized KV to the paged KV cache. The **Fused Fetch-and-Dequant** kernel loads quantized KV from HBM and dequantizes immediately into registers for reuse scenarios such as chunk prefill or prefix cache operations.

The memory subsystem is co-optimized through **cache-aligned tiling**. The tile size along the content dimension is increased from \(64\) to \(128\) in order to align with 128-byte L2 cache lines, match Swizzle-128B SMEM layout, and generate coalesced TMA descriptors. This is paired with **zero-overhead layout transformation** strategies: V-tile transposition is performed through an SMEM \(\rightarrow\) RF \(\rightarrow\) SMEM route, and QK accumulator layouts are permuted at byte granularity in registers to match PV input requirements. These transformations are overlapped with other compute segments so that their cost is largely hidden by the pipeline [2602.10718].

In effect, SnapMLA treats long-context decoding as a dataflow problem as much as a numerical one. The content cache is stored in FP8, the positional branch remains in BF16, scales are carried in forms compatible with online normalization, and the memory layouts are chosen to satisfy both Tensor Core and cache-line constraints. The result is an MLA-specific decoding stack rather than a generic adaptation of standard MHA quantization practice.

## 6. Empirical performance, relation to adjacent work, and limits

The reported evaluation uses **DeepSeek-V3.1**, a \(671\)B-parameter MoE model with \(37\)B active parameters, and **LongCat-Flash-Thinking**, a \(560\)B MoE model with \(18.6\)–\(31.3\)B active parameters. The baseline is **FlashMLA** in BF16. Benchmarks include MMLU-Pro, MMLU-Redux, IFEval, Arena-Hard, MATH-500, AIME-24/25, BeyondAIME, GPQA-Diamond, ZebraLogic, and LiveCodeBench [2602.10718].

The quality changes reported for SnapMLA are small and mixed in sign. For **DeepSeek-V3.1**, MMLU-Pro changes from \(84.41\) to \(84.43\), MMLU-Redux from \(90.48\) to \(90.89\), IFEval from \(86.32\) to \(87.25\), MATH-500 from \(98.80\) to \(98.20\), AIME-24 from \(93.85\) to \(93.65\), BeyondAIME from \(71.80\) to \(69.90\), GPQA-Diamond from \(84.15\) to \(82.57\), and LiveCodeBench from \(73.46\) to \(72.74\). For **LongCat-Flash-Thinking**, MMLU-Pro changes from \(82.60\) to \(82.73\), MMLU-Redux from \(89.30\) to \(88.27\), IFEval from \(86.90\) to \(87.80\), MATH-500 from \(99.20\) to \(98.80\), AIME-24 from \(91.87\) to \(91.67\), BeyondAIME from \(69.50\) to \(70.20\), and LiveCodeBench from \(79.40\) to \(79.74\). The paper characterizes these fluctuations as typically within roughly one point and uses them to support the claim that long-context reasoning and code-generation performance is preserved to near-BF16 levels.

Throughput gains are more pronounced. Across DP/TP configurations of DP1/TP8, DP4/TP2, and DP8/TP1 at context lengths from \(16\)k to \(128\)k, SnapMLA achieves up to a **\(1.91\times\)** throughput improvement over BF16 FlashMLA, with the largest gains in **DP8/TP1**, where FP8 memory savings permit larger batch sizes. At the kernel level, the paper reports an effective BF16 peak of \(148\) TFLOPS and derives an effective FP8 peak of approximately \(279.6\) TFLOPS from the MLA tile mix; SnapMLA is reported to track this effective peak closely across sequence lengths. With head counts \(H \in \{16,32,64,128\}\), performance increases with \(H\) and saturates around \(H \ge 64\) at roughly \(85\%\) of the effective peak. Multi-token prediction with query length \(2\) yields a modest additional boost [2602.10718].

Relative to adjacent systems, SnapMLA differs from **FlashAttention-3** in being MLA-specific rather than standard-MHA-specific, particularly because it must handle decoupled RoPE and shared latent KV structure. Relative to **FlashMLA**, it extends a BF16 Hopper kernel to FP8 through RoPE-aware quantization, pre-scaled QK, scale-fused PV, and MLA-specific fused operators. Relative to **TransMLA**, it addresses inference efficiency rather than architectural conversion: TransMLA converts GQA checkpoints into MLA structure, whereas SnapMLA assumes MLA structure and optimizes its decoding path [2502.07864].

The main limits are explicit. SnapMLA concerns **decoding**, not training; it is tied to **Hopper-class FP8 hardware**; and its methods are specific to **MLA**, not to arbitrary compressed-attention or state-space alternatives. The released implementation is integrated with **SGLang-FluentLLM** and is intended to replace BF16 FlashMLA kernels in the serving path. Within those bounds, SnapMLA represents a specialized solution to the long-context execution problem posed by MLA-based large language models [2602.10718].

Source: https://www.emergentmind.com/topics/snapmla