SnapMLA: FP8 Decoding for MLA Models

• SnapMLA is an FP8 decoding framework for MLA models that optimizes long-context inference by addressing quantization and system-level challenges.
• It introduces RoPE-aware per-token KV quantization and scale-fused PV computation to achieve up to 1.91× throughput improvements without significant quality loss.
• Leveraging hardware-aware kernel co-optimization on NVIDIA Hopper, SnapMLA restructures decoding dataflow for efficient FP8 inference in deep language models.

SnapMLA is an FP8 decoding framework for long-context inference in DeepSeek-style Multi-head Latent Attention (MLA) LLMs. 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 (Zhang et al., 11 Feb 2026).

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 (Zhang et al., 11 Feb 2026).

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 (Zhang et al., 2023). 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 (Meng et al., 11 Feb 2025). 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 (Zhang et al., 11 Feb 2026).

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 $h_t \in \mathbb{R}^d$0 is

$h_t \in \mathbb{R}^d$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 $h_t \in \mathbb{R}^d$2, 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 (Zhang et al., 11 Feb 2026).

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 (Zhang et al., 11 Feb 2026).

The first obstacle is heterogeneous quantization sensitivity. In the reported analysis for LongCat-Flash-Thinking, content components derived from $h_t \in \mathbb{R}^d$3 are tightly concentrated around zero, approximately within $h_t \in \mathbb{R}^d$4, whereas RoPE components $h_t \in \mathbb{R}^d$5 have a much wider dynamic range, up to $h_t \in \mathbb{R}^d$6, 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 $h_t \in \mathbb{R}^d$7. 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 (Zhang et al., 11 Feb 2026).

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, $h_t \in \mathbb{R}^d$8, and does so at per-token granularity, while keeping the RoPE component $h_t \in \mathbb{R}^d$9 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 (Zhang et al., 11 Feb 2026).

Per-token quantization is introduced in the standard form

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

with one scale $$c_{KV} = W^{DKV} h_t,\quad c_{KV} \in \mathbb{R}^{d_c},\ d_c \ll d.$$1 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 $$c_{KV} = W^{DKV} h_t,\quad c_{KV} \in \mathbb{R}^{d_c},\ d_c \ll d.$$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:

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

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

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

then the output

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

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

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

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 $$c_{KV} = W^{DKV} h_t,\quad c_{KV} \in \mathbb{R}^{d_c},\ d_c \ll d.$$7, SnapMLA applies block-wise dynamic quantization to $$c_{KV} = W^{DKV} h_t,\quad c_{KV} \in \mathbb{R}^{d_c},\ d_c \ll d.$$8, with block size $$c_{KV} = W^{DKV} h_t,\quad c_{KV} \in \mathbb{R}^{d_c},\ d_c \ll d.$$9, 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 (Zhang et al., 11 Feb 2026).

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 (Zhang et al., 11 Feb 2026). 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 $$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}.$$0 to $$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}.$$1 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 $$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}.$$2 RF $$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}.$$3 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 (Zhang et al., 11 Feb 2026).

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 $$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}.$$4B-parameter MoE model with $$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}.$$5B active parameters, and LongCat-Flash-Thinking, a $$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}.$$6B MoE model with $$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}.$$7–$$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}.$$8B 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 (Zhang et al., 11 Feb 2026).

The quality changes reported for SnapMLA are small and mixed in sign. For DeepSeek-V3.1, MMLU-Pro changes from $$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}.$$9 to $c_{KV}$0, MMLU-Redux from $c_{KV}$1 to $c_{KV}$2, IFEval from $c_{KV}$3 to $c_{KV}$4, MATH-500 from $c_{KV}$5 to $c_{KV}$6, AIME-24 from $c_{KV}$7 to $c_{KV}$8, BeyondAIME from $c_{KV}$9 to $k^R$0, GPQA-Diamond from $k^R$1 to $k^R$2, and LiveCodeBench from $k^R$3 to $k^R$4. For LongCat-Flash-Thinking, MMLU-Pro changes from $k^R$5 to $k^R$6, MMLU-Redux from $k^R$7 to $k^R$8, IFEval from $k^R$9 to $k_i^C$0, MATH-500 from $k_i^C$1 to $k_i^C$2, AIME-24 from $k_i^C$3 to $k_i^C$4, BeyondAIME from $k_i^C$5 to $k_i^C$6, and LiveCodeBench from $k_i^C$7 to $k_i^C$8. 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 $k_i^C$9k to $v_i^C$0k, SnapMLA achieves up to a $v_i^C$1 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 $v_i^C$2 TFLOPS and derives an effective FP8 peak of approximately $v_i^C$3 TFLOPS from the MLA tile mix; SnapMLA is reported to track this effective peak closely across sequence lengths. With head counts $v_i^C$4, performance increases with $v_i^C$5 and saturates around $v_i^C$6 at roughly $v_i^C$7 of the effective peak. Multi-token prediction with query length $v_i^C$8 yields a modest additional boost (Zhang et al., 11 Feb 2026).

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 (Meng et al., 11 Feb 2025).

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 LLMs (Zhang et al., 11 Feb 2026).