Tensor Layout Conventions

Everything in this entry is a coin flip that someone made once and that a few hundred million parameters were then fitted around. There is no argument to be had about which is correct; there is only the question of which one the checkpoint in front of you assumed.

Axis order

Two orderings for the attention tensors are in wide use.

Both describe the same numbers. The second keeps the head and feature dimensions adjacent in memory, so a token’s full activation is contiguous — which is what a fused kernel wants, and why the transpose between the two shows up in every profile of a naïve implementation.

RoPE pairing

Rotary embeddings rotate the feature dimension in pairs, and there are two ways to decide which components are paired.

The original paper specifies the first. Meta’s LLaMA reference implementation and therefore most of the ecosystem use the second, and Hugging Face’s conversion scripts permute the query and key weight matrices at load time so that the halved implementation reproduces the interleaved model.

That permutation is the part that bites. Load a checkpoint through a path that skips it and the model still runs, still produces English, and has an attention mechanism whose positional signal is scrambled — coherent for a clause, adrift by the end of the paragraph.

Epsilon placement

Both prevent the division by zero. They differ in the gradient near small $\sigma$, and they differ in the forward value by an amount that is negligible unless the activation is small — which, in a well-normalised network, some of them are. Every reference implementation puts it inside. A surprising number of reimplementations do not.

import torch
from torch import Tensor


def interleaved_to_halved(w: Tensor, heads: int, dim: int) -> Tensor:
    """Permute a q or k projection so a halved-RoPE kernel reproduces
    an interleaved-RoPE checkpoint. Applied to weights, once, at load."""
    return (
        w.view(heads, dim // 2, 2, -1)   # split each head's rows into pairs
        .transpose(1, 2)                 # pairs → (first halves, second halves)
        .reshape(w.shape)
    )


def rotate_halved(x: Tensor) -> Tensor:
    """The rotation this pairing implies: [-x2, x1] over the two halves."""
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)


def rotate_interleaved(x: Tensor) -> Tensor:
    """The other one. Same operation, different notion of 'pair'."""
    x = x.view(*x.shape[:-1], -1, 2)
    x1, x2 = x.unbind(-1)
    return torch.stack((-x2, x1), dim=-1).flatten(-2)

When a converted model is fluent and wrong, work through these before anything else, because each is a single line and each produces exactly that symptom:

• RoPE pairing — is the weight permutation applied, and applied to $q$ and $k$ only?
• Axis order — does the kernel expect [B, H, T, D] or [B, T, H, D]?
• Gate order in SwiGLU — the fused up-projection is a single matrix whose halves are gate and value; swapping them is silent.
• Norm variance — biased ($/d$) or Bessel-corrected ($/(d{-}1)$)?
• Attention scale — $1/\sqrt{d_{\text{head}}}$, not $1/\sqrt{d_{\text{model}}}$.
• Causal mask offset — does position $i$ see position $i$, or only $i-1$?

The fastest way to localise it is to compare hidden states layer by layer against a known-good implementation on one short prompt. The first layer where the cosine similarity drops below about 0.99 is the layer with the bug, and the list above will name it.