RoPE and FlashAttention: relative positions and online softmax

Two changes at different levels

RoPE changes how positional information enters attention scores. FlashAttention changes how a specified attention operation is computed and moved through memory. One is a model convention, the other an IO-aware implementation. They compose, but neither guarantees useful extrapolation or inexpensive arbitrarily long context.

Bring scaled dot-product attention, positional embeddings and the cache contract from DL038–042. By the end, derive the rotary identity, implement offset-aware rotations, reproduce tiled softmax attention on CPU, explain backward recomputation and choose an SDPA backend with a measured execution boundary. We do not install CUDA kernels or run a GPU benchmark here.

1 · RoPE from a relative-position requirement

Let q and k be projected content vectors in one two-dimensional plane. Rotate q by angle mθ and k by nθ:

Orthogonality gives R(a).T = R(-a) and angle addition gives the identity. For fixed q and k, the positional part depends on relative offset n−m. The score still depends on content; across layers the content vectors can themselves encode context. “Only distance matters” is false. Rotations preserve vector norms but do not guarantee that scores decay monotonically with distance: they oscillate.

For even head dimension d, rotate pairs (0,1),(2,3),... at frequencies theta_i = base**(-2*i/d). Standard RoPE rotates Q and K, not V; values are combined by the resulting weights. A different architecture could transform values too, but it would not be the conventional score-only operation demonstrated here.

Additive sinusoidal positions do not have the same generic projected-content identity: expanding (x+p)Wq against (y+r)Wk gives content-position cross terms depending separately on absolute positions. That does not establish that every additive-position model fails outside training length or that every rotary model succeeds.

2 · Offset-aware RoPE and meaningful tests

The original same-position elementwise test checked only diagonal q/k scores. Those correspond to zero relative offset and can miss bugs. Test all pairs at unequal positions, norm preservation and cached chunks. Keep interleaved and half-split conventions aligned with the checkpoint's Q/K projection row ordering; both can represent equivalent functions after a consistent permutation, but swapping the rotation alone changes a trained model.

import math
import numpy as np
import torch
import torch.nn.functional as F
from scipy.special import softmax
 
torch.set_num_threads(1)
torch.manual_seed(43)
 
def rope(x, positions, base=10000.0):
    if x.shape[-1] % 2 or base <= 0:
        raise ValueError('positive base and even head dimension required')
    if positions.ndim != 1 or positions.numel() != x.shape[-2]:
        raise ValueError('one absolute position per query/key row')
    d = x.shape[-1]
    inv = base ** (-torch.arange(0, d, 2, device=x.device,
                                 dtype=torch.float64) / d)
    angles = positions.to(device=x.device, dtype=torch.float64)[:, None] * inv
    c, s = angles.cos().to(x.dtype), angles.sin().to(x.dtype)
    a, b = x[..., 0::2], x[..., 1::2]
    return torch.stack([a*c-b*s, a*s+b*c], -1).flatten(-2)
 
q = torch.randn(2, 3, 7, 8, dtype=torch.float64)
k = torch.randn(2, 3, 9, 8, dtype=torch.float64)
qp, kp = torch.arange(4, 11), torch.arange(2, 11)
a = rope(q, qp) @ rope(k, kp).transpose(-2, -1)
b = rope(q, qp+23) @ rope(k, kp+23).transpose(-2, -1)
torch.testing.assert_close(a, b, rtol=1e-11, atol=1e-11)
torch.testing.assert_close(rope(q, qp).square().sum(-1), q.square().sum(-1))
wrong = rope(q, qp) @ rope(k, kp+1).transpose(-2, -1)
assert not torch.allclose(a, wrong)
full = rope(k, kp)
chunked = torch.cat([rope(k[..., :4, :], kp[:4]),
                     rope(k[..., 4:, :], kp[4:])], -2)
torch.testing.assert_close(full, chunked)
assert not torch.allclose(full[..., 4:, :], rope(k[..., 4:, :], torch.arange(5)))
print('RoPE all-pairs relative-shift, norm, chunk and wrong-offset controls pass')

Inside a decoder, project Q/K/V, compute absolute positions from cache position, rotate new Q/K, append rotated K and unrotated V, then run attention with the offset mask from DL042. Do not rotate cached K again. Remove the additive positional embedding if the intended architecture replaces it; adding both is another model. Cache cos/sin tables as device-aware buffers or compute slices as above. For long positions, low-precision angle construction can introduce errors even before attention, so keep the frequency/angle policy explicit.

Practice solution: with q=(1,0), k=(1,0), θ=1, the score at m=2,n=5 is cos(3). Shifting both by10 leaves cos(3); shifting only k gives cos(13). If your code checks only m=n, every correct orthogonal rotation gives q·k and you have not tested relative offsets.

3 · Context extension is not automatic

Increasing the RoPE base decreases the frequencies for nonzero pair indices; shrinking the base increases them. The original recommendation to shrink the base to keep angles small was reversed. The first frequency remains1 under base changes, so base scaling is not identical to dividing every position by a constant.

• Position Interpolation: map positions into the trained range, commonly by dividing positions by an extension factor, then continue training under the changed convention. This trades resolution for range; it is not an accuracy guarantee at every distance.
• NTK-aware base scaling: change frequency scales in a dimension-dependent way through a larger base. A particular checkpoint needs the matching formula and configuration, not an arbitrary number substituted at serving time.
• YaRN: combines frequency-dependent interpolation/extrapolation with an attention-temperature adjustment. Its paper studies efficient extension in particular models; “the default everywhere” is not an algorithmic property.
• LongRoPE: searches nonuniform rescaling factors and uses progressive extension/readjustment in its study. Its reported multi-million-token experiments are not proof that any Phi-3-mini checkpoint supports two million tokens out of the box.
• Llama3: the report lists RoPE base500,000 and separate long-context training. A larger base by itself does not transform an8k checkpoint into a validated128k model. Meta's April2025 Llama4 announcement describes interleaving layers without positional embeddings with rotary layers (iRoPE) and inference-time attention-temperature scaling. It says Scout was trained with256k context and advertises10M-token inference. That is a separate model/training recipe, not simply sliding-window layers or a consequence of Llama3's base value.

Evaluate short and long contexts, retrieval at varied positions, distractors and task accuracy, not merely whether allocation succeeds. An advertised window is a supported input size, not proof of uniform recall. Long input also costs tokenization, prefill, KV storage and decode bandwidth.

4 · The attention memory problem

Naive attention computes scores S=QKᵀ/sqrt(d), probabilities P=softmax(S) and O=PV. With batch1,32 heads,8192 tokens and two-byte scores, one full score tensor holds4,294,967,296 bytes, exactly4GiB. Scores, probabilities, gradients and other activations can multiply the allocation. This is distinct from the persistent inference KV cache.

HBM holds large tensors; registers and shared memory on each streaming multiprocessor hold much smaller working sets. Capacities and bandwidths depend on the device; “a few MB SRAM per SM” and fixed40GB HBM are not universal specifications. The goal is to avoid repeated HBM writes/reads of an intermediate that will be immediately reduced.

Picture a small notepad beside a warehouse: calculate one score tile, fold its contribution into a summary, then reuse the notepad. You still visit the required Q/K pairs; you do not need to retain their entire grid. Causal kernels may skip masked future pairs, so “every square-grid pair is always evaluated” is also too strong.

5 · Online softmax: sufficient running state

For one query row maintain a maximum m, exponential denominator l and unnormalized weighted-value numerator u. After a tile with scores s and values V:

Return u/l after the last tile. Rescale both accumulators when the maximum changes. If you instead store a normalized previous output O, multiply it by the old l before merging; mixing the two conventions is a common error. Start with m at negative infinity and zero l/u, and explicitly handle a tile with no valid keys so negative-infinity subtraction does not produce NaN. Define what an entirely masked row means; this example rejects it.

Here a Q-block outer loop retains row state across K/V blocks and writes final output once. That is an explanatory loop order, not a claim that FlashAttention-1's K/V-outer Algorithm1 writes every output only once. Different FlashAttention versions partition work differently.

def tiled_attention(q, k, v, allowed, block=3):
    q, k, v = [np.asarray(x, dtype=np.float64) for x in (q, k, v)]
    if block <= 0 or q.ndim != 2 or k.ndim != 2 or v.ndim != 2:
        raise ValueError('positive block and matrix inputs required')
    if q.shape[1] != k.shape[1] or len(k) != len(v):
        raise ValueError('Q/K feature and K/V length mismatch')
    allowed = np.asarray(allowed, dtype=bool)
    if allowed.shape != (len(q), len(k)) or not allowed.any(axis=1).all():
        raise ValueError('each query must have at least one valid key')
    out = np.empty((len(q), v.shape[1]))
    for i in range(0, len(q), block):
        qq = q[i:i+block]
        m = np.full(len(qq), -np.inf)
        den = np.zeros(len(qq)); num = np.zeros((len(qq), v.shape[1]))
        for j in range(0, len(k), block):
            scores = qq @ k[j:j+block].T / math.sqrt(q.shape[1])
            mask = allowed[i:i+block, j:j+block]
            scores = np.where(mask, scores, -np.inf)
            active = mask.any(axis=1)
            new = np.maximum(m[active], scores[active].max(axis=1))
            scale = np.exp(m[active] - new)
            weights = np.exp(scores[active] - new[:, None])
            den[active] = den[active]*scale + weights.sum(axis=1)
            num[active] = num[active]*scale[:, None] + weights @ v[j:j+block]
            m[active] = new
        out[i:i+block] = num / den[:, None]
    return out
 
rng = np.random.default_rng(43)
Q, K, V = rng.normal(size=(7, 4)), rng.normal(size=(9, 4)), rng.normal(size=(9, 5))
mask = np.arange(9)[None, :] <= (2 + np.arange(7))[:, None]
scores = np.where(mask, Q @ K.T / 2, -np.inf)
reference = softmax(scores, axis=-1) @ V
for block in (1, 2, 3, 20):
    np.testing.assert_allclose(tiled_attention(Q, K, V, mask, block), reference,
                               rtol=1e-12, atol=1e-12)
qt, kt, vt = [torch.tensor(x, dtype=torch.float64, requires_grad=True) for x in (Q,K,V)]
y = F.scaled_dot_product_attention(qt, kt, vt, attn_mask=torch.tensor(mask), dropout_p=0.0)
np.testing.assert_allclose(y.detach().numpy(), reference, rtol=1e-12, atol=1e-12)
y.square().sum().backward()
eps = 1e-6
qp, qm = Q.copy(), Q.copy(); qp[2,1] += eps; qm[2,1] -= eps
fd = (np.square(tiled_attention(qp,K,V,mask)).sum() -
      np.square(tiled_attention(qm,K,V,mask)).sum()) / (2*eps)
assert abs(fd-qt.grad[2,1].item()) < 1e-7
try:
    tiled_attention(Q,K,V,np.zeros_like(mask))
    raise AssertionError('empty row accepted')
except ValueError:
    pass
# A rising maximum requires rescaling earlier contributions.
x = np.array([0., 1., 20.]); values = np.array([3., 4., 9.])
naive_tile_average = (softmax(x[:2]) @ values[:2] + values[2]) / 2
assert abs(naive_tile_average - softmax(x) @ values) > 1
print('tiled attention, SciPy/torch parity, gradient and missing-rescale control pass')

This NumPy implementation teaches the algebra, not GPU IO performance. With fixed widths and tile size, persistent Q/K/V/O and row statistics use linear sequence storage; scratch tiles are bounded. However, this teaching harness also stores a dense allowed mask and materializes dense scores for its independent oracle, so the whole harness still uses quadratic memory. A linear-storage implementation must generate causal/padding eligibility per tile or use a compact mask representation, and omit the dense oracle from the measured path. Dense attention arithmetic remains quadratic. Backward can reconstruct probabilities from Q/K and saved row normalization rather than retain the full probability matrix, spending extra FLOPs to reduce activation storage.

“Exact attention” means the same real-arithmetic operation, not bit-for-bit equality across tile order, hardware or dtype. Floating-point reassociation, FP8 quantization and nondeterministic kernels require explicit tolerances and quality evaluation. Similar loss curves alone cannot prove two algorithms identical.

6 · SDPA dispatch and an unexecuted GPU comparison

For a square, unpadded causal self-attention call in torch2.9, F.scaled_dot_product_attention(q,k,v,dropout_p=0.0,is_causal=True) is a good interface. It is not a promise of FlashAttention on every Ampere+ input. Device, dtype, dimensions, mask and enabled backends constrain dispatch; supported explicit masks do not universally disable fused attention. SDPA applies the supplied dropout probability regardless of module eval mode, so pass zero explicitly for evaluation. For cached rectangular queries use the correct explicit offset mask and is_causal=False; never combine both flags under this API contract.

Need attention weights for inspection or an unsupported custom bias? A materialized reference can be appropriate. Benchmark small shapes too: launch overhead or unsupported layouts can reverse an assumed speedup. ALiBi/bias support is backend/version-dependent, not a permanent prohibition. Do not replace a custom attention operation with SDPA merely because its shapes fit.

The following complete callable procedure is syntax/source reviewed, not executed. It requires an existing compatible CUDA runtime and enough memory; no download or allocation is performed unless a reader explicitly calls it. It forces each backend, synchronizes, records peak allocation and compares outputs. Start below4096 and report unsupported/OOM cases rather than inventing4–6x gains.

# GPU reference: unexecuted; call only in an already approved CUDA environment.
def benchmark_sdpa_gpu(lengths=(128, 512, 1024), repeats=10):
    import time, statistics, torch
    import torch.nn.functional as F
    from torch.nn.attention import sdpa_kernel, SDPBackend
    if not torch.cuda.is_available():
        raise RuntimeError('CUDA required; no automatic provisioning')
    if repeats <= 0:
        raise ValueError('positive repeats required')
    results = []
    for length in lengths:
        torch.manual_seed(43)
        q = k = v = reference = None
        try:
            q,k,v = [torch.randn(2,8,length,64,device='cuda',dtype=torch.float16)
                     for _ in range(3)]
            with torch.inference_mode(), sdpa_kernel(SDPBackend.MATH):
                reference = F.scaled_dot_product_attention(q,k,v,dropout_p=0.0,is_causal=True)
        except RuntimeError as exc:
            results.append(dict(length=length,backend='reference-setup',error=str(exc)))
            del q,k,v,reference
            continue
        for backend in (SDPBackend.MATH, SDPBackend.FLASH_ATTENTION,
                        SDPBackend.EFFICIENT_ATTENTION):
            try:
                with torch.inference_mode(), sdpa_kernel(backend):
                    for _ in range(3):
                        out = F.scaled_dot_product_attention(q,k,v,dropout_p=0.0,is_causal=True)
                    torch.cuda.synchronize()
                    error = (out-reference).abs().max().item()
                    baseline = torch.cuda.memory_allocated()
                    torch.cuda.reset_peak_memory_stats()
                    samples = []
                    for _ in range(repeats):
                        torch.cuda.synchronize(); t = time.perf_counter()
                        out = F.scaled_dot_product_attention(q,k,v,dropout_p=0.0,is_causal=True)
                        torch.cuda.synchronize(); samples.append(time.perf_counter()-t)
                    results.append(dict(length=length,backend=str(backend),
                        median_seconds=statistics.median(samples),max_error=error,
                        peak_allocated=torch.cuda.max_memory_allocated(),
                        incremental_peak=torch.cuda.max_memory_allocated()-baseline))
            except RuntimeError as exc:
                results.append(dict(length=length,backend=str(backend),error=str(exc)))
        del q,k,v,reference
    return dict(torch=torch.__version__,device=torch.cuda.get_device_name(),results=results)

7 · Successors, MLA and other long-context levers

FA1 introduced IO-aware exact attention; FA2 reduces non-matmul work and improves work partitioning/parallelism; FA3 exploits Hopper asynchrony, warp specialization and low precision. The official2024 FA3 account separates roughly740TFLOP/s FP16 (about75% utilization) from about1.2PFLOP/s FP8: the original combined75% and1.3PFLOP/s as if they were one measurement. These are paper/blog hardware results, not local tests. Private vendor adoption, alleged FA3+ universal Blackwell kernels and dollar savings for a whole training run cannot be inferred from a kernel benchmark.

MLA and decoupled RoPE: if expanded key is W_UK*c, inserting a position-dependent rotation between the query and this projection generally prevents absorbing a single fixed W_UK into the query. V3 splits score features into a nonrotary compressed path and a separate small rotary Q/K path. Cache latent c and the rotary key; the nonrotary projections can be absorbed while positions remain explicit. FlashMLA supplies specialized kernels for supported MLA layouts/hardware. The main README captured on2026-10-08 announces a2026-09-30 compatibility break for earlier models/Hopper and points users of the older layout to commitba89a3466e9470ad08ab39738d4e7bb66989e1e7; that pinned README separately identifies V3/V3.2 kernels. Thus the original2025 V3 discussion must not silently use the latest main branch. Neither README's performance supports “V3 is5–10x cheaper because of this kernel.”

Ring Attention: distribute sequence blocks and circulate K/V, overlapping communication with local block attention. This addresses per-device storage and work distribution; total dense pair arithmetic remains quadratic and network bandwidth matters. “Near-infinite context” is not a finite-cluster capacity bound, and the paper does not establish an undisclosed Gemini implementation.

Window/sparse attention: omit selected pairs, changing connectivity; multi-layer receptive fields and retained global tokens determine accessible dependencies. Linear attention / Performer: reorganize kernelized or approximated attention through feature maps; the softmax-equivalence claim no longer applies in general. Mamba/SSMs: use structured recurrence/selective state rather than the same dense softmax operation. These are distinct representational choices, not further FlashAttention versions.

Choose by measured bottleneck: fused dense attention for expensive score materialization; GQA/KV quantization for cache bytes; positional extension for a model's usable range; sparsity/recurrence for pair-compute limits; sequence parallelism when interconnect and distributed capacity justify it. RoPE rescaling, quantization and architecture changes need quality tests, not only timing. Even exact fused kernels need numerical acceptance.

8 · Recall, practice and complete answers

1. State the RoPE identity. R(mθ)ᵀR(nθ)=R((n−m)θ), so fixed content vectors' positional interaction depends on offset. It does not remove their content dependence or prove extrapolation.
2. Reconstruct apply-RoPE. Split even/odd coordinates; return interleaved (a*cos-b*sin, a*sin+b*cos) using absolute position angles and one frequency per pair. Standard V is unchanged. The program above supplies all dimension/device/offset details.
3. Diagnose convention mismatch and rotated V. Compare intermediate Q/K after a known row permutation; align projection weights as well as rotations. If weights match but the value path is rotated, the weighted content sum changed even if scores match. Do not blame cache allocation for that discrepancy.
4. Why linear storage but quadratic work? Keep Q/K/V/O plus row statistics and a tile, consume every permitted score without retaining the full grid, and generate the mask per tile. The teaching harness's dense mask and independent oracle do not satisfy that total-memory bound. Backward recomputes scores; no claim of linear dense arithmetic follows.
5. What is online softmax? Maintain max, denominator and numerator; rescale old sums when max rises, add new contributions and normalize at the end. Averaging separately normalized tile outputs is wrong because their partition functions differ.
6. Does one API call guarantee FlashAttention? No. SDPA dispatch is capability-dependent. Force a backend to test eligibility, set eval dropout to zero and use the proper mask. Explicit attention weights or unsupported operations may require a reference path.
7. Why decouple RoPE in MLA? Position-dependent rotation cannot generally commute through a fixed low-rank up-projection. Keep the compressible nonrotary latent separate from a small rotary key/query component. This retains a cache-efficient form without discarding positional score information.
8. Stretch solution: implement the recurrence from section5 and compare against SciPy softmax for tile sizes1,2,3 and larger than sequence length. Include all-masked tiles, a rising maximum and a rectangular causal mask. The executed tests do so and compare an input derivative with torch autograd. This is algebra verification, not a reproduced GPU speedup.

A90-second explanation should distinguish algorithm, storage, traffic and arithmetic. Tomorrow derive the recurrence without looking; after a week repeat the all-pairs offset test and explain why its diagonal-only predecessor was insufficient.

Sources and optional navigation

RoFormer, FlashAttention, FA2, official FA3 explanation, torch2.9 SDPA, PI, YaRN, LongRoPE, Ring Attention, DeepSeek-V3, FlashMLA. Exact primary passages and dated raw snapshots accompany the review.

Optional inherited videos: RoPE, FlashAttention; YaRN explainer. These are not evidence for the removed benchmark promises. Course map · Previous DL042 · Next DL044.