Attention Through Arithmetic Intensity

jxmorris121 pts0 comments

[Tech] Why MLA and MTP Fight Each Other: Attention Through Arithmetic Intensity | ChangyiYang's Site

Table of Contents

I was reading this post by Su Jianlin when one sentence stopped me:

Besides the KV cache, decoding now has another variable — MTP, or speculative decoding, whose idea is to trade compute for speed. But MLA behaves during decoding like an MQA with head_dims=512+, and has already consumed most of the compute up front, so “MLA+MTP” tends to lose out.

My first reaction was: what does that mean? Why would MLA “consume compute up front” during decode? And why should it conflict with MTP that its decode FLOPs come out equivalent to a head-dim-512+ MHA?

Following the question down led to something classic, but unusually pretty when applied to attention: arithmetic intensity — FLOPs per byte moved.

And the answer turns out to be surprisingly clean (everything below assumes a BF16 KV cache):

Reduce MHA all the way down and its AI comes out to exactly 1 ;

GQA and MQA are just as clean — they depend on nothing but head counts . Context length and head dim cancel out completely;

MLA has the same shape again , independent of the latent dim too, just with a constant of a little under 2 in front.

Lined up, here is the AI of the attention core for a single-token decode:

Attention<br>what the cache holds<br>AI, roughly

MHA<br>each query head has its own K and V

GQA<br>a group of query heads shares one K and V<br>query heads / KV heads

MQA<br>all query heads share one K and V<br>number of query heads

MLA<br>one latent, K and V both expand from it<br>~2 × number of query heads

Sections 2 and 4 derive those four rows. The first three are the same formula with different KV head counts; the constant on the last row comes from somewhere else entirely.

And that constant of a little under 2 is just enough to move attention decode on many current GPUs from clearly memory-bound to sitting near the roofline knee. Stack MTP on top and the workload tips over into compute-bound — which is exactly why Su says MLA is unfriendly to MTP.

This post works through the whole derivation.

This one is written out in detail — starting from how you count FLOPs in a matmul, with every matrix shape spelled out. If you already know the structure of attention and what decode computes, section 2 can be skimmed down to the result in 2.5, then jump to section 3.

1. What arithmetic intensity is#

The definition is simple:

That is: for every byte pulled in from HBM, how many floating-point operations do you get out of it.

Low AI means the GPU spends most of its time moving data. High AI means each piece of data gets reused for a lot of arithmetic once it arrives.

The hardware has a matching threshold:

In the idealized roofline model:

and

AI_{\text{hardware}} \quad\Rightarrow\quad \text{compute-bound}" loading="lazy">

Using dense BF16 tensor-core throughput, the theoretical balance points of a few common cards:

GPU<br>Dense BF16 Peak<br>HBM Bandwidth<br>Theoretical Balance Point

H100 SXM<br>~989.5 TFLOP/s<br>3.35 TB/s<br>~295 FLOP/B

H200 SXM<br>~989.5 TFLOP/s<br>4.8 TB/s<br>~206 FLOP/B

B200 (HGX)<br>~2.25 PFLOP/s<br>~8 TB/s<br>~281 FLOP/B

Sources: NVIDIA H100, NVIDIA H200, NVIDIA HGX B200, NVIDIA DGX B200.

No real kernel saturates peak FLOPs and peak bandwidth simultaneously, so treat these as a roofline upper bound for building intuition rather than a line you would see in a profiler. &ldquo;A few hundred FLOP/B&rdquo; is the magnitude to remember; it gets compared against later.

One caveat worth stating up front: AI is a ratio. It answers &ldquo;which side of the roofline are you on&rdquo;, not &ldquo;which approach is faster.&rdquo; Numerator and denominator can both grow and leave AI untouched while everything gets slower. Section 5 has a concrete case: two algorithms computing the same thing, where the one with the higher AI does 120× the FLOPs.

What follows only counts the KV-related part of attention:

For one decoded token: from its hidden state, compute Q, K and V, read in the KV cache, and carry through to this layer&rsquo;s attention output.

Softmax is small next to the two big matmuls and is left out. The goal is not to estimate whole-layer latency; it is to isolate one question: what does changing the attention structure do to the arithmetic intensity of that stretch?

Sections 2 through 4 handle decode only (one token at a time, history read from cache). Prefill — computing an entire sequence at once — waits until section 5, where it turns out the same model wants the opposite algorithm in the two phases.

PS: why not count the final WO too? Because it has nothing to do with which attention structure you picked. WO always receives the concatenated per-head outputs, whose width depends only on Hq and dv — how the KV side is organized is invisible to it, and MHA, GQA, MQA and MLA all hand it something the same width. Like the Q/K/V projections it is a weight-times-vector: it never touches the KV cache, does not grow with L, and at...

attention from arithmetic cache compute decode

Related Articles