Softmax plus cross-entropy gets summarized so often as p - y that it is easy to forget where that expression comes from:

The formula is right, but two implementation details are easy to miss.

First, softmax is not elementwise. Every output probability depends on every logit in the same row because the classes share one normalization denominator. So the backward pass is a Jacobian-vector product, not a coordinate-wise derivative.

Second, \(p-y\) is a gradient with respect to logits, not probabilities. It shows up only after the log in cross-entropy interacts with the softmax Jacobian.

This note walks through both pieces: the general softmax backward pass for an arbitrary upstream gradient, and the cross-entropy/logsoftmax case that produces the expression used in training code.

Setup

For one row, let the logits be

and define

For a batch, write the logits as

and the softmax outputs as

Softmax is applied row-wise:

Here \(i\) indexes the batch row. The indices \(j\) and \(k\) index classes, vocabulary entries, or whatever lives on the final axis. Rows do not interact, so it is enough to derive the formula for one row and apply it to the whole batch.

The Softmax Derivative

For one batch row \(i\), define the row-wise denominator

Then

We want the derivative

This asks a local question: if we nudge the logit for class \(k\), what happens to the softmax probability assigned to class \(j\)?

Using the quotient rule,

The numerator derivative depends on whether \(j=k\):

where

The denominator derivative is

Substituting both pieces gives

Since

we get the scalar softmax derivative:

This compact expression contains both cases:

and, for \(j \neq k\),

The sign tells the story. Increasing a logit’s own value increases its own probability. Increasing some other logit gives that other class more of the shared denominator, so this probability goes down.

The Jacobian

For a fixed row \(i\), collect all derivatives into a Jacobian matrix

The diagonal entries are \(S_{ij}(1-S_{ij})\). The off-diagonal entries are \(-S_{ij}S_{ik}\). In matrix form:

because

and

So each entry is

This Jacobian is dense, but it is also highly structured. We definitely do not want to materialize a \(K \times K\) matrix for every row. For a language model vocabulary, \(K\) can be tens or hundreds of thousands. The backward pass needs the Jacobian-vector product, not the Jacobian itself.

There is a clean reason this Jacobian is symmetric. Define

Its gradient is softmax:

So the softmax Jacobian is the Hessian of logsumexp:

For a smooth scalar function, the Hessian is symmetric. That is why \(J_i^\top=J_i\) in this case.

Backpropagating Through Softmax

Let the upstream gradient for one row be

For a vector function \(y=f(x)\), the usual column-vector backpropagation rule is

For softmax, \(J_i\) is symmetric, so

Now expand the multiplication:

The first term is elementwise multiplication:

The second term has a scalar dot product inside:

Therefore,

Factoring out \(S_i\) coordinate-wise:

Elementwise:

This computes the same thing as the full Jacobian multiply, but the actual work is just a row-wise dot product and an elementwise multiply:

def softmax_backward(grad_out, softmax_out):
    dot = (softmax_out * grad_out).sum(dim=-1, keepdim=True)
    return softmax_out * (grad_out - dot)

Here grad_out is \(\partial L / \partial S\), and softmax_out is the saved softmax output \(S\) from the forward pass. The keepdim=True keeps the dot product shaped as \(B \times 1\) so it broadcasts across the class dimension.

As a quick sanity check, each row of \(G_Z\) sums to zero. That has to happen: adding the same constant to every logit in a row does not change softmax, so the backward pass should not push in that all-ones direction.

Cross-Entropy With Logits

Now let the target be \(y\). For a one-hot target, exactly one coordinate is \(1\). For label smoothing or soft targets, assume \(y\) is a normalized target distribution:

The cross-entropy loss for one example is

where

Since

we can rewrite the loss directly as a function of logits:

Because \(\sum_j y_j=1\),

Differentiate with respect to logit \(z_i\):

and

Therefore,

and in vector form,

For the target class \(t\), this is

For a non-target class \(i \neq t\), it is

In code, the plain batch-mean version looks like this:

probs = torch.softmax(logits, dim=-1)
B = logits.shape[0]
dlogits = probs.clone()
dlogits[torch.arange(B, device=logits.device), targets] -= 1
dlogits /= B

The final dlogits tensor is not a probability distribution. It is the gradient \(\partial L/\partial Z\). For a batch mean reduction with no class weights and no ignored labels,

The Same Result Through the Softmax Jacobian

The logsumexp derivation is the shortest route. It is still useful to check the same result by chaining cross-entropy through the softmax Jacobian.

Starting from

the gradient with respect to probabilities is

elementwise.

The softmax backward formula from above is

Compute the scalar dot product:

Then

This is the cancellation people usually have in mind. The dense softmax Jacobian is still there mathematically, but the \(1/p_i\) term from differentiating the log collapses it into a simple logits gradient.

LogSoftmax Backward

In practice, frameworks usually avoid computing softmax probabilities and then taking a log. They use logsoftmax or a fused cross-entropy kernel, because logsumexp can be computed stably by subtracting the row maximum before exponentiating.

Define

Differentiate:

Let

Then

So the vector form is

and the batch form is

Here rowsum has shape \(B \times 1\) and broadcasts across the final axis.

In code:

def logsoftmax_backward(grad_out, softmax_out):
    rowsum = grad_out.sum(dim=-1, keepdim=True)
    return grad_out - softmax_out * rowsum

For negative log-likelihood or cross-entropy,

so

Because

substituting into the logsoftmax backward pass gives

At that point, all three routes agree:

and, for a plain batch mean,

A Small Numerical Check

Suppose the predicted probabilities and one-hot target are

Then

The target-class gradient is negative. Under gradient descent,

subtracting a negative value increases the target logit. The non-target gradients are positive, so their logits decrease. That matches the training signal we expected.

What to Remember

Softmax backward is a structured Jacobian-vector product:

Cross-entropy with logits simplifies to:

with whatever scaling the loss reduction adds.

The mental model I use is this: softmax couples the classes through a shared denominator, and cross-entropy contributes the reciprocal probability term that removes that coupling in the logits gradient. That is why the training code can use a simple \(P-Y\) update even though softmax itself has a dense Jacobian.

Follow me on Twitter, Github or connect on LinkedIn.

The idea of Decompose-K and the custom-op autotuning workflow comes from the PyTorch Conference talk Lightning Talk: Faster Than SOTA Kernels in Torch.compile With Subgraph Fusions and Custom Op Autotuning - Elias Ellison & Paul Zhang, Meta. This post is my own implementation walkthrough and benchmark study built around that idea.

The skinny large-K matmul problem

A standard matmul is

C[M, N] = A[M, K] @ B[K, N]

and the way a GPU GEMM extracts parallelism is by tiling the M x N output. Each program owns a BLOCK_M x BLOCK_N tile of C and streams over K to accumulate it. That works well when M and N are large, because there are many output tiles and the GPU has plenty of independent work to fill its streaming multiprocessors (SMs).

The problem case is a skinny, K-dominant matmul: M and N are tiny while K is huge. Think M = N = 16 with K = 32768, or a decode-time MoE router GEMM like [T, 7168] @ [7168, 256] where T can be as small as 1. Now the output is 16 x 16 = 256 elements, which is one or two tiles. The GPU has 132 SMs sitting idle while one or two programs serially walk a reduction of length 32768. The matmul is reduction-bound, but the standard tiling exposes almost no parallelism along the only large axis.

Decompose-K is a restructuring that fixes exactly this mismatch. The basic idea is simple: if the only big dimension is K, then split K and parallelize over the split.

What Decompose-K does

Split the K dimension into S independent chunks, compute S partial GEMMs, then sum the partials:

A[M, K] @ B[K, N]
  -> partials[S, M, N]
  -> sum(partials, dim=0)

Each partial is a smaller matmul over K/S of the reduction. The S partials are independent, so they become a batched matmul (bmm) with batch dimension S. The minimal PyTorch version makes this concrete:

def decomposeK(a, b, k_splits):
    m, k = a.shape
    n = b.shape[1]
    assert k % k_splits == 0, "k must be divisible by k_splits"
    k_parts = k // k_splits

    # [m, k] -> [m, k_splits, k_parts] -> [k_splits, m, k_parts]
    a_reshaped = a.reshape(m, k_splits, k_parts).permute(1, 0, 2)
    b_reshaped = b.reshape(k_splits, k_parts, n)        # [k_splits, k_parts, n]

    result = torch.bmm(a_reshaped, b_reshaped, out_dtype=torch.float32)
    reduced_result = result.sum(dim=0)
    return reduced_result.to(a.dtype)

The important part is what the reshape buys. For M = N = 16, K = 32768, S = 64:

• a_reshaped is [64, 16, 512], b_reshaped is [64, 512, 16].
• The bmm now has 64 independent matmuls instead of one. That is 64 units of work the scheduler can spread across SMs, versus a single output tile before.
• Each partial accumulates only 512 of the reduction, not 32768.

We have traded one long serial reduction for S short parallel ones, plus a final reduction of the S partials. The partials are accumulated in fp32 (out_dtype=torch.float32) so the split does not cost accuracy relative to a single fp32-accumulated matmul.

This is essentially split-K, but expressed at the tensor level as a bmm plus a reduction rather than as atomic adds into a single output tile. That distinction matters once we add an epilogue, which is the next point.

Why it is epilogue-friendly

A split-K design that uses atomic adds into the output has a problem if you want to fuse an elementwise epilogue like ReLU: the output tile is not final until every split has finished its atomic contribution, so you cannot apply ReLU during the accumulation. You would need a separate pass after all atomics settle.

Decompose-K keeps the partials in a separate [S, M, N] buffer and does an explicit reduction. That means the reduction step is the natural and only place where each output element becomes final, so an epilogue can be folded directly into the reduction’s store:

acc = sum over splits of partials[:, m, n]
acc = relu(acc)          # fused into the same kernel
store C[m, n] = acc

No extra pointwise pass over C, no second read/write of the output. For tiny outputs that are memory-bound on the epilogue, this is a real saving, and we will measure it later (~1.2x–1.4x over an unfused ReLU).

Where it is worth it

Decompose-K is attractive when:

• K is very large and M/N are small (e.g. M = N = 16..64, K = 8192..32768).
• The workload is latency-sensitive and a single fixed shape matters more than a general GEMM. A concrete example is a DeepSeek-V3 MoE router GEMM [T, 7168] @ [7168, 256], where decode has tiny dynamic T = 1..256 and prefill has larger T.
• A fused epilogue like ReLU can ride along on the reduction.

It is not worth it when M and N are already large enough to fill the GPU, when K is small, when K divides poorly for the candidate split counts, or when the extra [S, M, N] buffer and its reduction dominate the cost.

The rest of this post is a tour of implementations of this one idea, from the laziest (torch.compile) to a hand-written Triton kernel that beats Inductor’s own autotuned choice. Every benchmark below is BF16 on an H100 (132 SMs), over the grid M = N ∈ {16, 32, 48, 64} and K ∈ {8192, …, 32768}.

Baseline: just call torch.compile

The first thing to try is to write decomposeK in plain PyTorch and let Inductor handle the rest. The relevant detail is the compile mode. Across the three benchmark suites used throughout this post - a BF16 matmul with a fused ReLU epilogue (epilogue-bf16), a plain BF16 matmul (matmul-bf16), and a plain FP32 matmul - max-autotune-no-cudagraphs was the best mode, edging out max-autotune:

decomposeK_compiled = torch.compile(decomposeK, mode="max-autotune-no-cudagraphs")

max-autotune turns on Inductor’s template autotuning (it benchmarks several generated kernels and picks the fastest). The -no-cudagraphs variant skips CUDA graph capture, which for these tiny single-shot calls avoids capture overhead without losing the autotuning benefit.

What does naive compilation actually emit?

Compiling the decomposeK function above (for the router shape [64, 7168] @ [7168, 256], S = 4) produces two operations, which you can read off the Inductor output code:

# extern bmm into an fp32 partials buffer
buf0 = empty_strided_cuda((4, 64, 256), (16384, 256, 1), torch.float32)
extern_kernels.bmm_dtype(
    reinterpret_tensor(arg0_1, (4, 64, 1792), ...),
    reinterpret_tensor(arg1_1, (4, 1792, 256), ...),
    out_dtype=torch.float32, out=buf0)

# one generated pointwise kernel: sum over the 4 splits + cast to bf16
triton_poi_fused__to_copy_sum_0.run(buf0, buf1, 16384, ...)

So the bmm goes to an external (cuBLAS) batched kernel, and the sum(dim=0) plus the .to(bf16) cast get fused into a single generated Triton pointwise kernel. The generated reduction kernel is literally an unrolled add of the S = 4 slices:

tmp0 = tl.load(in_ptr0 + (x0), None)
tmp1 = tl.load(in_ptr0 + (16384 + x0), None)
tmp3 = tl.load(in_ptr0 + (32768 + x0), None)
tmp5 = tl.load(in_ptr0 + (49152 + x0), None)
tmp7 = (tmp0 + tmp1 + tmp3 + tmp5)
tl.store(out_ptr0 + (x0), tmp7, None)

This is the call graph for explicit Decompose-K written in PyTorch: bmm + a fused sum/cast kernel. If instead you write the version with a ReLU epilogue, the fused kernel additionally folds in maximum(0, x), so you get bmm + a fused sum+relu kernel. The epilogue is free in the sense that it rides on the reduction kernel that has to run anyway.

What if you just write relu(mm(a, b)) and let Inductor decide?

This is the more interesting question, because the PyTorch nightly used here (torch==2.12.0.dev20260408+cu128) ships a Decompose-K lowering inside Inductor itself - see torch/_inductor/template_heuristics/decompose_k.py and the subgraph choice it registers in torch/_inductor/kernel/mm.py. So Inductor will reach for the decomposition on large-K shapes on its own, autotuning it as one more candidate against the regular matmul templates. The POC compiles a plain torch.relu(torch.mm(a, b)) and dumps the generated code at two K values.

Small K (M = N = 16, K = 256) - Inductor emits a single fused matmul template, triton_tem_fused_mm_relu_0, with the source nodes [aten.mm, aten.relu]. The ReLU is fused into the matmul template’s store suffix:

# inductor's template suffix, inside the matmul kernel
tmp1 = triton_helpers.maximum(tmp0, acc)   # relu
tl.store(out_ptr1 + xindex, tmp1, mask)

One kernel, ReLU fused, done. There is no reason to decompose at small K.

Large K (M = N = 16, K = 32768) - Now Inductor chooses Decompose-K on its own. The generated graph is named decompose_k_mm_64_split_5 (it picked S = 64, so k_part = 512) and contains three pieces:

# 1) batched partial matmul via cuBLAS, fp32 accumulate
extern_kernels.bmm_dtype(
    reinterpret_tensor(arg0_1, (64, 16, 512), ...),
    reinterpret_tensor(arg1_1, (64, 512, 16), ...),
    out_dtype=torch.float32, out=buf0)         # buf0: [64, 16, 16] fp32

# 2) generated reduction over the 64 splits
triton_per_fused_mm_0.run(buf0, buf2, 256, 64, ...)

# 3) a SEPARATE pointwise relu kernel
triton_poi_fused_relu_1.run(buf1, 256, ...)

The thing to notice is piece 3. When Inductor takes the Decompose-K lowering, it emits ReLU as a separate triton_poi_fused_relu_1 pointwise kernel after the reduction. It does not fuse ReLU into the Decompose-K reduction/store. That is an extra full read-and-write of the output buffer. For a tiny 16 x 16 output this is small in absolute terms, but it is exactly the fusion opportunity a hand-written kernel can reclaim, and it is the gap the rest of this post chases.

So we have two facts to build on: Decompose-K is the right structure at large K (Inductor agrees), and the stock Inductor lowering leaves the epilogue unfused. Time to write the kernel ourselves.

A hand-written Triton kernel

Source: kernels/decompose_k_triton_kernel.py

The kernel is two stages that mirror the structure above: a partial-matmul kernel that fills [S, M, N], and a reduction/epilogue kernel that sums over S and optionally applies ReLU on the store.

Stage 1: partial matmul

The partial-matmul kernel uses a 2D launch grid: program_id(0) indexes the M x N output tile (with the usual L2-friendly group-major swizzle), and program_id(1) indexes the split. Each program computes one BLOCK_M x BLOCK_N tile for one split, accumulating only its K // SPLIT_K slice of the reduction.

@triton.jit
def _partial_mm(a, b, partials, ...):
    pid = tl.program_id(0)
    split_id = tl.program_id(1)

    # group-major swizzle of pid -> (pid_m, pid_n) for L2 reuse
    ...
    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    offs_k = tl.arange(0, BLOCK_K)

    k_per_split = K // SPLIT_K
    split_k_start = split_id * k_per_split
    acc = tl.zeros((BLOCK_M, BLOCK_N), tl.float32)

    for k0 in range(0, k_per_split, BLOCK_K):
        k_offsets = k0 + offs_k
        a_ptrs = a + offs_m[:, None] * stride_am + (split_k_start + k_offsets[None, :]) * stride_ak
        b_ptrs = b + (split_k_start + k_offsets[:, None]) * stride_bk + offs_n[None, :] * stride_bn
        k_mask = k_offsets < k_per_split
        a_vals = tl.load(a_ptrs, mask=(offs_m[:, None] < M) & k_mask[None, :], other=0.0)
        b_vals = tl.load(b_ptrs, mask=k_mask[:, None] & (offs_n[None, :] < N), other=0.0)
        acc += tl.dot(a_vals, b_vals, out_dtype=tl.float32, input_precision=INPUT_PRECISION)

    partial_ptrs = partials + split_id * stride_ps + offs_m[:, None] * stride_pm + offs_n[None, :] * stride_pn
    tl.store(partial_ptrs, acc, mask=(offs_m[:, None] < M) & (offs_n[None, :] < N))

A few details worth calling out:

• The accumulator is fp32 regardless of input dtype, and input_precision is "ieee" for fp32 inputs and "tf32" otherwise. This keeps the split from changing numerical behaviour versus a single accumulated matmul.
• split_k_start = split_id * k_per_split is the only thing that distinguishes one split program from another. Each split reads a contiguous k_per_split band of K.
• The store writes into the split-indexed partials[split_id] slice. There are no atomics: every (split_id, tile) pair owns a disjoint region of the partials buffer.

Stage 2: reduce + fused epilogue

The reducer launches one program per output tile, loops over all SPLIT_K partials into a tile-shaped accumulator, applies ReLU if requested, and stores:

@triton.jit
def _reduce_epilogue(partials, c, ..., SPLIT_K, BLOCK_M, BLOCK_N, FUSE_RELU):
    # (pid -> pid_m, pid_n) swizzle as before
    ...
    acc = tl.zeros((BLOCK_M, BLOCK_N), tl.float32)
    for split_id in range(0, SPLIT_K):
        acc += tl.load(ptrs + split_id * stride_ps,
                       mask=(offs_m[:, None] < M) & (offs_n[None, :] < N), other=0.0)

    if FUSE_RELU:
        acc = tl.maximum(acc, 0.0)

    tl.store(c_ptrs, acc, mask=(offs_m[:, None] < M) & (offs_n[None, :] < N))

This is the fusion that Inductor’s Decompose-K lowering does not do: ReLU is applied in registers before the single store of C, with no separate pointwise pass. Correctness-wise, this is safe precisely because the explicit reduction is where each output element first becomes final.

This kernel is correct and reasonable, but it carries a structural limitation in how the reducer is parallelized - one we will pin down in a moment. Because the surprising part comes first: this hand-written kernel does not actually beat Inductor.

Custom-op autotuning: letting Inductor pick the decomposition

Source: custom_op_autotune_relu_dispatch.py

Inductor exposes an API, register_custom_op_autotuning, that lets you hand it a list of alternative decompositions for an op and have it benchmark and select among them per shape, then lower the winner. The neat trick is that the target op can be either a real @torch.library.custom_op or an existing ATen overload like torch.ops.aten.mm.default. So you can intercept the lowering of every torch.mm in a compiled graph.

The candidates are plain mm (or mm + relu) plus one Decompose-K decomposition per valid split count:

K_SPLITS = (2, 4, 8, 16, 32, 64, 128, 256)

def generate_mm_relu_configs(fake_tensors):
    k = int(fake_tensors["a"].shape[1])
    splits = [s for s in K_SPLITS if k % s == 0]
    configs = [CustomOpConfig(mm_relu_impl)]
    configs += [CustomOpConfig(decompose_k_relu_impl, k_splits=s) for s in splits]
    return configs

decompose_k_relu_impl is just the PyTorch-level bmm + sum + relu from the start of the post; we are not handing Inductor a Triton kernel here. We are handing it several mathematically-equivalent PyTorch decompositions and letting it lower and time each one.

The script registers this at two different boundaries, each with a matching config generator, so both the plain-matmul and the fused matmul+ReLU cases are covered:

• aten.mm boundary - generate_mm_configs, keyed self/mat2. Candidates are mm_impl (the ordinary torch.mm lowering) plus decompose_k_impl(k_splits=s) for every s in K_SPLITS that divides K. ReLU stays outside the autotuned op as a separate pointwise kernel.
• fused mm_relu custom-op boundary - generate_mm_relu_configs, keyed a/b. Candidates are mm_relu_impl (plain relu(mm)) plus decompose_k_relu_impl(k_splits=s) for each valid s - the ReLU is folded into every candidate, so Inductor times the fully fused decomposition directly.

The shape of the candidate list is the same either way: the plain lowering plus one Decompose-K candidate per valid split count. Decompose-K is never special-cased - it is just one more entry in the menu that has to win the autotune on its own merits at each shape.

The registration wires up the autotuning:

register_custom_op_autotuning(
    custom_op=torch.ops.aten.mm.default,    # intercept every torch.mm
    config_generator=generate_mm_configs,   # candidates from fake-tensor shapes
    name="router_mm_relu_autotune",
    input_gen_fns={                          # make real CUDA tensors to benchmark
        "self": lambda fake: torch.randn_like(fake, device="cuda") * 0.1,
        "mat2": lambda fake: torch.randn_like(fake, device="cuda") * 0.1,
    },
    dispatch_on={"tensor_name": "self", "dim": 0, "range_upper_bound": 1024},
    split_points=[1, 8, 32, 128, 512],
    benchmark_with_cudagraphs=True,
)

The pieces that matter:

• config_generator receives fake tensors keyed by the op’s schema argument names (self/mat2 for aten.mm) and returns candidates for the current compile shape. This is how it inspects K and only emits splits where K % k_splits == 0.
• input_gen_fns turn the fake tensors into real CUDA tensors so each candidate can actually be timed. The keys must match the ATen schema names.
• dispatch_on + split_points enable range-based dispatch. Here it benchmarks and dispatches on self.shape[0], i.e. the M/T dimension. split_points = [1, 8, 32, 128, 512] becomes ranges roughly [1,1], [2,8], [9,32], [33,128], [129,512], [513, inf]. Inductor picks a winner per range and, if adjacent ranges want different winners, emits a runtime torch.cond dispatch tree. This is the “custom op dispatching per dynamic shape” idea: decode (T=1) and prefill (T large) can get different kernels from the same compiled graph.

So there are two layers of timing in the benchmark. First, Inductor autotunes during compile and lowers a winner. The harness captures that decision into the CSV: across the whole grid, the winner was decompose_k_relu_impl with k_splits of 64 or 128 (it agrees that Decompose-K is right). Second, the harness times the already-compiled callable with do_bench.

Two ways to use the registration: one graph for all shapes, or one per shape

The registration is used in two different ways.

Standalone, custom_op_autotune_relu_dispatch.py runs the dynamic path: it registers with dispatch_on + split_points, compiles with dynamic=True, and sweeps T = [1, 16, 64, 256, 768] against eager. Here the torch.cond per-M dispatch tree matters - one compiled graph routes decode-like (T = 1) and prefill-like (T large) shapes to whatever candidate won their range.

The benchmark measures one fixed shape at a time, so it skips that path. bench_decompose_k.py resets Dynamo (torch._dynamo.reset()) before each grid point and compiles with dynamic=False, so Inductor specializes for the exact (M, K, N). It registers the static variants (register_mm_relu_static_autotune / register_mm_static_autotune), which leave out dispatch_on and split_points: with a single shape there is no range to cover and no torch.cond tree, so each shape is autotuned and lowered on its own and the harness records the winner. Range-based dispatch belongs to the standalone exploration; it plays no part in the Results numbers below.

An easy way to mis-benchmark: the Dynamo recompile limit

This grid is easy to mis-benchmark. torch.compile specializes the same Python function over many exact shapes, and TorchDynamo’s default config.recompile_limit is 8 per code object. If you keep recompiling for new K without resetting Dynamo, you eventually hit:

torch._dynamo hit config.recompile_limit (8)

After that, later shapes stop getting fresh optimized graphs and silently fall back to slower execution, which makes the recorded custom-op timing (the custom_op_mm_relu_ms column in the benchmark CSV) look great for the first few shapes and then jump up to the eager band. That is a benchmark cache artifact, not the candidate getting slower. The fix is to reset Dynamo between exact-shape grid points before compiling the next measured callable (compile time is not part of the latency measurement, so this is fair).

Limitations of the hand-written kernel

How does our hand-written baseline Triton kernel compare against this custom-op autotuned path? It loses, everywhere:

Speedup is custom_op_mm_relu_ms / decompose_k_fused_ms, using the column names in the benchmark CSV - the standalone kernel is logged as decompose_k_fused_ms or decompose_k_unfused_ms depending on whether ReLU is fused into the reduction, and as decompose_k_ms in the plain matmul-bf16 suite. Below 1.0x means the standalone kernel is slower. Inductor’s lowering of the same Decompose-K math beats our kernel by ~8–13% on the median. The reason is the reducer - the structural limitation hinted at earlier.

Our reducer is output-tile shaped: one program owns a BLOCK_M x BLOCK_N tile, carries a 2D accumulator of that shape, and serially walks the split dimension in a Python-level for loop. For a tiny output like M = N = 16, a 16 x 16 reducer tile can mean a single reducer program for the whole output, and that one program serially reads all SPLIT_K partials. The reduction parallelism is tied to the matmul output tiling, which is the wrong axis to parallelize when the output is tiny and the split count is large.

Inductor sidesteps exactly this: its generated reduction (a triton_per_fused persistent reduction over the split axis) parallelizes the split correctly. The matmul template is roughly a wash; the reduction is where we are losing - and that is precisely what the optimized kernel fixes next.

The optimized Triton kernel

Source: kernels/decompose_k_triton_kernel_optimized.py

The optimized kernel keeps the same two-stage structure (and reuses the exact same _partial_mm) but rewrites the reducer and widens the autotuning search. There are four changes.

1. Reshape the reducer around the split axis

This is the big one. Instead of an output-tile-shaped accumulator that serially loops over splits, flatten the output matrix into a 1D element index x = m * N + n and treat the split as the reduction axis of a proper 2D vector reduction:

vals: [XBLOCK, RBLOCK]   # XBLOCK output elements x RBLOCK splits
acc:  [XBLOCK]           # one fp32 result per output element

Each reducer program owns a disjoint slice of XBLOCK flattened output elements, loads all RBLOCK (= SPLIT_K, padded to a power of two) partials for them, and reduces with a single tl.sum over the split axis:

@triton.jit
def _reduce_epilogue_vector_flat(partials, c, stride_ps, XNUMEL, SPLIT_K, XBLOCK, RBLOCK, FUSE_RELU):
    x_base = tl.program_id(0) * XBLOCK + tl.arange(0, XBLOCK)
    x = x_base[:, None]
    r = tl.arange(0, RBLOCK)[None, :]
    vals = tl.load(partials + r * stride_ps + x,
                   mask=(x < XNUMEL) & (r < SPLIT_K), other=0.0)
    acc = tl.sum(vals, 1)                  # vector reduction over splits
    if FUSE_RELU:
        acc = tl.maximum(acc, 0.0)
    tl.store(c + x_base, acc, mask=x_base < XNUMEL)

For M = N = 16 (256 output elements) with XBLOCK = 32, this launches 8 reducer programs, each owning 32 elements and reducing all splits for them in one vectorized sum. There is no cross-program combine step: every program writes final values directly. Crucially, this decouples reducer tiling from matmul tiling: the partial matmul can use 16x16 or 64x64 tiles while the reducer independently uses a small XBLOCK chosen for split-reduction efficiency.

2. A flat contiguous fast path

For the common case of contiguous row-major tensors, the address math collapses. Because partials[s, m, n] lives at base + s*(M*N) + (m*N + n) and x = m*N + n, the hot loads and stores are just partials + r * stride_ps + x and c + x_base - no division/modulo to recover (m, n):

if partials.is_contiguous() and c.is_contiguous():
    _reduce_epilogue_vector_flat[grid](...)   # no (m,n) reconstruction
else:
    _reduce_epilogue_vector[grid](...)        # general strided fallback

The strided kernel (with the x // N, x - m*N reconstruction) stays as a fallback for non-contiguous views like c = base[:, ::2].

3. Warp counts matched to tiny tiles

The baseline search used 4 warps everywhere. But a 16x16 output tile has only 256 fp32 accumulators; even with a K slice of 64–128, a partial-matmul program for that tile is too small to keep 4 warps (128 threads) busy. Four warps there means more scheduling/sync overhead, more register pressure, and fewer resident programs per SM - the opposite of what a tiny-output, large-K shape calls for.

So the optimized config set deliberately includes 1- and 2-warp small tiles (16x16x64, 16x16x128, …) alongside 4-warp larger tiles (64x32, 64x64, …). The intuition is work per program:

16x16x128 =  32,768 multiply-add positions   -> few warps
64x32x64  = 131,072 multiply-add positions
64x64x128 = 524,288 multiply-add positions   -> 4 warps pay off

Both families are in the search, and the benchmark picks per shape. (In the results, small M=N tend to pick 16x16 tiles, while M=N=64 shifts to 64x32/64x64.)

4. More split candidates

The optimized split search explicitly tries the power-of-two counts (2, 4, 8, 16, 32, 64, 128, 256) and then appends the baseline divisor-based candidates. This matches the custom-op setup and maps naturally onto the power-of-two RBLOCK the vectorized reducer uses for the split axis.

Now it wins

With the vectorized reducer, the standalone kernel flips from losing everywhere to winning almost everywhere against the custom-op autotuned path:

The reducer rewrite alone moved the median from 0.917x to 1.026x on the epilogue suite - roughly an 11% swing, recovered entirely from how the split reduction is parallelized and from folding ReLU into the store instead of running Inductor’s separate pointwise pass.

Benchmark setup

A few details about how the numbers below are produced, since they affect what the comparison actually measures.

One fresh compile per shape - The grid sweeps M = N ∈ {16, 32, 48, 64} against K ∈ {8192, 12288, 16384, 20480, 24576, 28672, 32768} - 28 shapes per suite. For every shape the harness calls torch._dynamo.reset() and then torch.compile(target, mode="max-autotune-no-cudagraphs", dynamic=False). The reset clears the compile cache so the recompile limit from earlier never triggers, and dynamic=False lets Inductor specialize fully for that one (M, N, K) instead of generating a shape-generic kernel. The cost is one autotuning pass per shape, but the payoff is that each point in the grid reflects the best kernel Inductor can produce for exactly that shape - which is the fairest thing to compare a hand-tuned kernel against. For FP32 runs the harness also sets torch.set_float32_matmul_precision("highest") (and "high" otherwise) so the matmul precision matches the chosen dtype.

Timing - Latency comes from triton.testing.do_bench(fn, warmup=10, rep=50, return_mode="median") - 10 warmup iterations, 50 measured, reported as the median. do_bench handles the L2-cache flush and CUDA-event timing internally, so the numbers are wall-clock kernel time without host overhead. Every candidate is also checked for correctness with torch.testing.assert_close (tolerances per dtype) before it is timed, so a config that diverges numerically is rejected rather than ranked.

Data types - The post shows two suites: epilogue-bf16 (a BF16 matmul with a fused ReLU epilogue) and matmul-bf16 (plain BF16 matmul, no epilogue). The repo runs two more that are not plotted here - matmul-fp16 (plain FP16 matmul) and matmul-fp32 (plain FP32 matmul, with tighter rtol=1e-4/atol=1e-3 tolerances). The FP16 picture matches BF16 closely; FP32 has less to gain because the larger element size and tighter accumulation leave less headroom for the split reduction, but the ordering is the same.

Results

All raw numbers and plots are checked into the repo, so the runs are reproducible and inspectable: the original hand-written kernel lives under bench_results/ and the optimized kernel under bench_results_v2/. Each suite has a .csv with per-shape timings (eager_ms, compiled_ms, custom_op_mm_relu_ms, decompose_k_fused_ms, decompose_k_unfused_ms, …) and the captured autotune winner, plus per-M=N and overall-grid plots.

The clearest way to see the whole picture is the overall comparison grid. Each x-axis point is a (M=N, K) shape; lower latency is better. Five curves: eager torch.mm + relu, compiled torch.mm + relu, the custom-op autotuned mm+relu, standalone Decompose-K with a separate ReLU, and standalone Decompose-K with fused ReLU.

First, the original standalone Triton kernel (red/purple sit in the middle of the pack, above the green custom-op line - the kernel is not yet competitive):

Then the optimized kernel. The fused Decompose-K curve (purple) drops to the bottom of the grid, at or below the green custom-op line across nearly all shapes:

Zooming into M = N = 16, the most reduction-bound slice, makes the ordering crisp: eager (blue) is slowest, compiled and separate-ReLU Decompose-K (orange/red) are in the middle, and the fused Decompose-K (purple) and custom-op (green) share the floor - with the fused kernel slightly ahead at small K:

The plain matmul suite (no epilogue) tells the same story with a smaller margin, since there is no epilogue to fuse - the win there is purely from the reducer:

A few representative numbers from the epilogue-bf16 runs - a BF16 matmul with a fused ReLU epilogue (do_bench median, ms):

Two things to read off this table. The fused Decompose-K kernel is consistently the fastest column, ~1.5–1.7x over eager and ~1.2–1.4x over compiled torch.mm + relu. And the last column - the fusion benefit alone, measured as the same Decompose-K config with ReLU as a separate in-place op versus fused into the reduction store - is a steady 1.19x–1.4x. That is the concrete payoff of the epilogue-friendliness from the very first section.

Takeaways

• Decompose-K is a structural fix for a structural problem. When the only large dimension is K, splitting it turns one long serial reduction into S parallel partial GEMMs plus a reduction, giving the GPU work to fill its SMs. It is most useful for skinny, large-K, latency-sensitive shapes (MoE routers, small-batch decode).
• It is epilogue-friendly by construction. Keeping partials in a separate buffer and doing an explicit reduction means the reduction’s store is the natural place to fuse ReLU - worth ~1.2–1.4x here. A split-K-with-atomics design cannot do this cleanly.
• torch.compile already knows the trick. At large K, Inductor picks Decompose-K on its own (extern bmm_dtype + a generated reduction). But it emits the epilogue as a separate pointwise kernel, leaving the fusion on the table.
• Custom-op autotuning is a strong, low-effort baseline. Handing Inductor a list of PyTorch decompositions and a per-range dispatch policy beat a naive hand-written Triton kernel on every shape. If you only do one thing, do this.
• Beating it required getting the reducer right. The win was not in the matmul; it was reshaping the reduction around the split axis (a real tl.sum over splits, decoupled from matmul tiling), matching warp counts to tiny tiles, and folding ReLU into the store. That moved the standalone kernel from 0/28 to 26/28 wins versus Inductor’s own choice.

It is worth being honest about the effort-to-reward ratio. The hand-written kernel wins, but only by a few percent over custom-op autotuning, and getting there meant rewriting the reducer and widening the search. For most situations, staying with torch.compile is perfectly reasonable - as long as it is done carefully. Plain torch.compile(decomposeK) was not enough on its own; Inductor decomposes but leaves the epilogue as a separate kernel, and it was that gap that pushed me toward custom-op autotuning. Set up that way, the compiler does a great job, and the hand-written kernel is the last few percent you reach for only when the shape is fixed and the latency genuinely matters.

And this gap is not fundamental. The day Inductor’s Decompose-K lowering learns to fuse the epilogue into the reduction store - the same optimization the hand-written kernel relies on - most of this margin trims away on its own, and the compiler path absorbs the win for free.

All the kernels, the custom-op autotuning setup, the benchmark harness, and the raw results are on GitHub: shreyansh26/MLSys-Experiments/decompose-k.

Follow me on Twitter, Github or connect on LinkedIn.

A useful map for KV cache work: are we selecting existing token states, or synthesizing new compact states? Is the cost paid per query, per corpus, or once as reusable training?

Code: cartridges, STILL-Towards-Infinite-Context-Windows, kv-cache-compression

Primary sources: Cartridges, STILL / neural KV-cache compaction, StreamingLLM / Attention Sinks, H2O, SnapKV, L2 norm KV compression

Long context is not only a modeling problem. It is also a memory residency problem.

During autoregressive decoding, the model does not recompute all previous tokens at every step. It stores their attention keys and values in the KV cache, then each new token attends to that cache. This is exactly what makes decoding practical. It is also what makes long prompts expensive to keep alive.

The basic tension is:

• the full KV cache is faithful, reusable, and expensive;
• a textual summary is cheap, but it has already passed through a narrow language bottleneck;
• retrieval is sparse and query-dependent;
• KV cache compression tries to keep the representation in the model’s own internal coordinate system.

This post is a code-first tour through that design space. I will start with the cache accounting, then walk through four token-selection style implementations: Attention Sink, L2 norm pruning, SnapKV, and H2O. After that, we will spend most of the time on the two more interesting compaction methods I implemented in detail: Cartridges and STILL.

The important distinction is that compression often means “keep fewer existing KV entries”, while compaction means “make a smaller KV object that is not just a subset of the original tokens.” Cartridges and STILL are compaction methods in that stronger sense.

Setup: The Memory Bill

Before comparing compression policies, it helps to make the cache cost explicit and name the questions each method has to answer.

The KV Cache Baseline

In a decoder-only transformer, the query for the current token attends to keys and values from previous tokens. For one layer and one head, ignoring masks for a moment:

where

The cache stores the $k_i$ and $v_i$ tensors so the model can append one new position at a time. With grouped-query attention, the number of query heads can be larger than the number of KV heads, but the storage formula is still simple:

Here $T$ is the number of cached tokens, $L$ is the number of layers, $H_{kv}$ is the number of KV heads, $d_{head}$ is the head dimension, the factor $2$ is for keys plus values, and $b$ is bytes per scalar.

The figure separates the growing $T$ term from the fixed per-token footprint. This post is mostly about replacing $T$ with a smaller cache budget $p$.

For a Llama-3.1-8B style configuration with $L=32$, $H_{kv}=8$, $d_{head}=128$, and bf16 storage, each token costs:

which is 128 KiB per token per request. A 128K-token prompt is therefore about 16 GiB of KV cache before allocator overheads, paging metadata, batching effects, and attention workspace. This is why “the model supports 128K” and “I can cheaply keep hundreds of 128K sessions resident” are very different statements.

The baseline cache update in the kv-cache-compression repo is exactly what you expect:

class KVCache(nn.Module):
    def __init__(self, max_batch_size, max_seq_length, n_heads, head_dim, dtype=torch.bfloat16):
        super().__init__()
        cache_shape = (max_batch_size, n_heads, max_seq_length, head_dim)
        self.register_buffer('k_cache', torch.zeros(cache_shape, dtype=dtype))
        self.register_buffer('v_cache', torch.zeros(cache_shape, dtype=dtype))

    def update(self, input_pos, k_val, v_val):
        assert input_pos.shape[0] == k_val.shape[2]

        k_out = self.k_cache
        v_out = self.v_cache
        k_out[:, :, input_pos] = k_val
        v_out[:, :, input_pos] = v_val

        return k_out, v_out

Every compression method below replaces this update policy, or replaces the whole cache object that the model consumes.

Four Questions For Any KV Compression Method

The implementation details vary, but I find the following questions more useful than method names:

1. What is retained? Original token positions, synthesized token positions, quantized vectors, or a mixture?
2. When is the decision made? During prefill, during every decode step, offline before serving, or through reusable training?
3. What signal decides importance? Recency, attention mass, key norms, distillation loss, or a learned encoder?
4. What does the model see afterward? A shorter prefix, a sliding cache, a trainable past-key-value object, or compact K/V plus attention biases?

Token Eviction: Keep A Subset Of The Original Cache

The small methods in kv-cache-compression are easy to understand because they keep the model architecture fixed and swap only the cache policy. Cartridges and STILL are more ambitious: they ask whether the cache itself can become a learned memory object.

Attention Sink: Keep The First Tokens And The Tail

StreamingLLM starts from a surprising empirical observation: pure sliding-window attention can collapse even when the recent local window is present. Keeping a small number of initial tokens fixes much of the instability. Those initial tokens act as “attention sinks”: places where many later tokens can put excess attention mass.

The policy is simple. Given a maximum cache budget

keep:

The paper’s key observation is not that the first few tokens are semantically important. It is that many models use the earliest positions as attention sinks: safe places to allocate attention mass when no specific old token is needed. Once a pure sliding window evicts those positions, the attention distribution shifts in a way the model was not trained for.

StreamingLLM's cache layout: the retained cache is not a contiguous suffix; it keeps initial sink tokens plus a tail window ending at the latest token.

In code, prefill with a long prompt keeps the first global_tokens and the last sliding_window tokens:

if total_len > self.max_cache_size:
    global_idxs = torch.arange(self.global_tokens, device=input_pos.device)
    recent_idxs = torch.arange(total_len - self.sliding_window, total_len, device=input_pos.device)
    keep_idxs = torch.cat([global_idxs, recent_idxs])
    new_pos = input_pos[keep_idxs]

    self.pos = new_pos.unsqueeze(0).unsqueeze(0).expand_as(self.pos)
    self.k_cache = k_val.index_select(dim=2, index=keep_idxs)
    self.v_cache = v_val.index_select(dim=2, index=keep_idxs)

During decoding, the implementation preserves the global region and evicts one position from the tail region:

idx_to_pop = (self.global_tokens + torch.argmin(self.pos[:, :, self.global_tokens :], dim=-1)).flatten()
self.pos[:, :, idx_to_pop] = input_pos.long()
self.k_cache[:, :, idx_to_pop] = k_val
self.v_cache[:, :, idx_to_pop] = v_val

The important part is what this method does not try to do. It does not identify facts. It does not reconstruct old values. It does not learn a memory. It keeps the model numerically stable in streaming use by preserving a few special early tokens and a recent local window.

That makes it a good default when the application is genuinely streaming: chat, logs, or continuous text where old details are less important than keeping the model coherent.

L2 Norm Compression: Score Keys Before Querying Them

The L2 norm strategy is interesting because it avoids attention statistics. The paper reports a correlation between key-vector norms and later attention behavior: low-norm keys tend to receive higher attention. That means we can score cached entries by the key tensor itself, before future queries arrive.

For a key vector $k_i$, define:

Then keep the tokens with the largest $s_i$, equivalently the lowest key norms.

Key logic from the L2 norm KV compression paper: score cached keys by norm and prune without building an observation attention matrix.

In the implementation, the score is written as max_norm - norm, then sorted:

key_norm = torch.norm(self.k_cache, p=2, dim=-1)
key_norm_diff = key_norm.max() - key_norm
scoring_priority = key_norm_diff.masked_fill(self.pos == -1, float('inf'))
scoring_sorted_idx = torch.argsort(scoring_priority, dim=-1)
num_toks_to_remove = int((1 - self.keep_ratio) * self.max_cache_size)
scoring_sorted_idx_selcted = scoring_sorted_idx[:, :, num_toks_to_remove:]

Because key_norm_diff is small for large-norm keys and large for small-norm keys, removing the first num_toks_to_remove sorted positions drops large-norm keys and keeps lower-norm keys. The implementation then gathers the retained K/V rows and appends empty slots for future decode tokens.

self.k_cache = torch.cat([
    torch.gather(
        self.k_cache,
        dim=2,
        index=scoring_sorted_idx_selcted.unsqueeze(-1).expand(-1, -1, -1, self.k_cache.shape[-1]),
    ),
    torch.zeros(
        self.k_cache.shape[0],
        self.k_cache.shape[1],
        num_toks_to_remove,
        self.k_cache.shape[-1],
        device=self.k_cache.device,
        dtype=self.k_cache.dtype,
    ),
], dim=2)

The useful property is that this is query-agnostic. The cache can be pruned without running an observation window or storing attention history. That also makes it compatible with attention kernels where the implementation does not expose full attention matrices.

The tradeoff is that the score is a proxy. It can work surprisingly well, but it is still betting that the norm structure learned by the model will line up with future importance.

SnapKV: Use Prompt Attention To Select The Cache

SnapKV is more query-aware, but it tries to make the decision before generation. The idea is that the prompt’s final tokens already reveal what the model will look for during decoding. So SnapKV observes attention from a small prompt window, scores old cache positions, keeps the top positions, and also keeps the most recent window.

The implementation path is:

1. Run prefill and keep the full prompt cache temporarily.
2. Compute attention scores from the prompt.
3. Aggregate query heads that share a KV head.
4. Sum attention from the last window_size query tokens to older prefix tokens.
5. Smooth the score with average pooling.
6. Keep the top compress_length - window_size older tokens plus the latest window_size tokens.

The “snap” in SnapKV is this one-time compression after prefill: the model observes the prompt, predicts which prefix positions will matter, and then generates against a much smaller cache.

SnapKV's selection step: the prompt observation window scores clustered prefix positions, then selected prefix K/V is concatenated with the observation window.

The core code is:

attn_scores_grouped = attn_scores.view(b, n_kv_head, n_rep, lq, lk)
attn_scores_agg = attn_scores_grouped.sum(dim=2)
attn_weights_sum = attn_scores_agg[:, :, -self.window_size:, : total_len-self.window_size].sum(dim=-2)
attn_cache = self.pool(attn_weights_sum)

indices = attn_cache.topk(self.compress_length - self.window_size, dim=-1).indices
indices_expanded = indices.unsqueeze(-1).expand(-1, -1, -1, head_dim)

k_past_compress = self.k_cache[:, :, :total_len-self.window_size, :].gather(dim=2, index=indices_expanded)
v_past_compress = self.v_cache[:, :, :total_len-self.window_size, :].gather(dim=2, index=indices_expanded)
k_cur = self.k_cache[:, :, total_len-self.window_size:total_len, :]
v_cur = self.v_cache[:, :, total_len-self.window_size:total_len, :]

key_states = torch.cat([k_past_compress, k_cur], dim=2)
value_states = torch.cat([v_past_compress, v_cur], dim=2)

The method is not “top attention globally.” It is “top attention from the observation window, smoothed locally, with a recent tail always retained.” That distinction matters because local neighborhoods can matter even when a single token’s raw attention score is not maximal.

In model_kv_cache_compression.py, SnapKV has to use a manual attention path during prefill so the cache object can see the scores:

y, scores = self.manual_attention(
    q,
    k,
    v,
    is_causal=(seqlen > 1),
    enable_gqa=(self.n_head != self.n_local_heads),
)

if self.snapkv_enabled and seqlen > 1:
     self.kv_cache.prune_cache(input_pos, scores, n_kv_head=self.n_local_heads)

That is the systems cost of attention-based pruning: the cache policy needs the attention matrix, and highly optimized attention kernels usually avoid materializing it.

H2O: Keep Heavy Hitters Under A Fixed Budget

H2O treats KV eviction as a dynamic heavy-hitter problem. Some tokens repeatedly receive a large share of attention. Those tokens should remain in cache, while low-utility tokens can be evicted.

For token $i$, maintain an attention-history score:

where $a_{\tau,i}$ is the attention mass assigned to token $i$ at decode step $\tau$, and $n_i(t)$ counts how many times the token was eligible or observed. When the cache is full, evict the smallest $h_i(t)$.

H2O also keeps recency in the picture. The useful mental model is not “attention history instead of a window”; it is “heavy hitters plus recent tokens under one fixed cache budget.”

H2O's eviction rule: accumulated attention identifies persistent heavy hitters, while recent tokens stay for local continuity.

The implementation stores numerator and denominator buffers:

self.register_buffer("attn_history_num", torch.zeros(history_num_shape, dtype=history_num_dtype))
self.register_buffer("attn_history_denom", torch.zeros(history_denom_shape, dtype=torch.int32))
self.register_buffer("attn_counter", torch.zeros((max_batch_size, n_heads), dtype=torch.int64))

The eviction score is the average attention history:

numerator = self.attn_history_num.sum(dim=-1).float()

if (self.history_window_size == 1):
    denominator = self.attn_history_denom.clamp_min(1)
else:
    denominator = self.attn_history_denom.clamp(min=1, max=self.history_window_size)

avg_attn = numerator / denominator
scores = avg_attn.masked_fill(self.pos == -1, float('-inf'))

During decode, once the cache is full, the least important token is replaced:

scores = self._calculate_eviction_scores()
eviction_idx = torch.argmin(scores, dim=-1)

eviction_idx_kv = eviction_idx.view(batch_size, n_heads, 1, 1).expand(-1, -1, 1, head_dim)
self.k_cache.scatter_(2, eviction_idx_kv, k_val)
self.v_cache.scatter_(2, eviction_idx_kv, v_val)

H2O is a good mental bridge between simple sliding windows and learned compaction. It says that recency alone is not enough; persistent attention mass is also important. But it still preserves original K/V entries. Once a token is gone, its vector is gone.

What Token Eviction Can And Cannot Do

Token-selection methods are useful because they are local changes to an inference stack. They keep the model frozen, keep the attention interface mostly unchanged, and usually require no offline training.

But they have a hard ceiling:

The compressed cache is a subset of original token states. It cannot store a new vector that mixes multiple facts. It cannot put “the answer to three distant questions” into one slot unless one original token state already encoded that mixture. It also tends to be query-shaped: SnapKV depends on the observed prompt; H2O depends on previous decode attention; Attention Sink depends on streaming recency.

Cartridges and STILL move to a different object:

where $p \ll T$, but $(k^c_j, v^c_j)$ need not equal any original token’s K/V pair. The slots can become learned memory.

Learned Cache Artifacts: Compaction Instead Of Eviction

The next methods do not merely choose which original token states survive. They build compact K/V artifacts whose slots can move away from individual token identity.

Cartridges: A Trainable KV Cache Per Corpus

The Cartridges paper asks a direct question: if users repeatedly ask questions against the same long corpus, can we train a compact KV cache for that corpus and reuse it?

The paper calls the training recipe self-study. First, the model uses the full corpus to generate supervision conversations about that corpus. Then the compact cache is trained with a context-distillation objective so future queries can use the cartridge instead of re-prefilling the whole source.

In the local implementation, a cartridge for one corpus stores, for every layer:

The base model is frozen. The trainable parameters are the compact K/V tensors themselves.

The object is initialized from a prefix pass:

@torch.no_grad()
def initialize_from_prefix_text(
    model,
    tokenizer,
    text: str,
    num_tokens: int,
    num_frozen_tokens: int = 1,
) -> TrainableKVCartridge:
    encoded = tokenizer(text, return_tensors="pt", add_special_tokens=False)
    input_ids = encoded["input_ids"][..., :num_tokens].to(model.device)
    outputs = model(input_ids=input_ids, use_cache=True)
    past_key_values = _normalize_past_key_values(outputs.past_key_values)
    keys = [layer[0].detach().to(model.dtype) for layer in past_key_values]
    values = [layer[1].detach().to(model.dtype) for layer in past_key_values]
    return TrainableKVCartridge(keys=keys, values=values, num_frozen_tokens=num_frozen_tokens)

This prefix initialization is easy to misread. It does not mean the final cartridge only contains the first $p$ tokens. It means optimization starts from the KV states of the first $p$ tokens. After training, the trainable slots are just parameters. They can move away from their original token identity.

The implementation also freezes a few initial positions as attention-sink-like anchors:

if num_frozen_tokens > 0:
    frozen_key = nn.Parameter(key_tensor[..., :num_frozen_tokens, :], requires_grad=False)
    frozen_value = nn.Parameter(value_tensor[..., :num_frozen_tokens, :], requires_grad=False)
    train_key = nn.Parameter(key_tensor[..., num_frozen_tokens:, :])
    train_value = nn.Parameter(value_tensor[..., num_frozen_tokens:, :])

The layer is materialized by concatenating frozen and trainable regions:

def layer(self, index: int) -> tuple[torch.Tensor, torch.Tensor]:
    key_parts = []
    value_parts = []
    if self.num_frozen_tokens > 0:
        key_parts.append(self.frozen_keys[index])
        value_parts.append(self.frozen_values[index])
    key_parts.append(self.trainable_keys[index])
    value_parts.append(self.trainable_values[index])
    return torch.cat(key_parts, dim=-2), torch.cat(value_parts, dim=-2)

At inference time, the cartridge is converted into the Hugging Face cache format:

def as_cache(self, model_config) -> DynamicCache:
    return DynamicCache(ddp_cache_data=self.as_legacy_past_key_values(), config=model_config)

So from the frozen model’s point of view, a cartridge is just past_key_values.

Cartridges as an offline compilation workflow: generate supervision from the full corpus, optimize a compact K/V artifact, then reuse it for follow-up queries.

The Cartridge Objective

The training signal in the local implementation is teacher distillation. First, a teacher answers using the full context. Then the cartridge is optimized so the frozen model, when given only the user query plus the cartridge cache, matches the teacher’s answer-token distribution.

For one answer position $t$, let $q_t(i)$ be the teacher distribution over a sparse set of candidate tokens. Let

be the frozen model’s probability for token $i$, given user query $u$ and cartridge $C_s$ for source corpus $s$. The loss is:

Concretely, _sparse_distillation_loss is the helper that implements the equation. It receives student logits already sliced to the assistant-token positions, then reconstructs a sparse teacher distribution from stored top-logprobs. The supervision argument is the per-token view of the teacher answer:

supervision: list[TokenSupervision]
  length == assistant_logits.shape[0]

TokenSupervision:
  token_id: int        # teacher's sampled/generated assistant token at this position
  logprob: float       # teacher log p(token_id)
  top_logprobs:
    - token_id: int    # candidate vocabulary id from the teacher's top-k distribution
      logprob: float   # teacher log p(candidate)
    - token_id: int
      logprob: float
    ...

So supervision[row_idx] describes the teacher distribution for the same assistant position as logits[row_idx]. In the raw synthesized conversation this information is stored as assistant token_ids plus sparse top_logprobs; the loss helper uses the expanded row-by-row form because it makes the alignment with the student logits explicit.

For example, suppose the teacher’s answer starts with two assistant tokens, roughly " Paris" and " is". Using illustrative token ids, the supervision might look like:

assistant_logits.shape = [2, vocab_size]

supervision =
  - token_id: 12041        # decoded token: " Paris"
    logprob: -0.10
    top_logprobs:
      - token_id: 12041    # " Paris"
        logprob: -0.10
      - token_id: 987      # " London"
        logprob: -2.40
      - token_id: 14321    # " Lyon"
        logprob: -3.10

  - token_id: 374          # decoded token: " is"
    logprob: -0.18
    top_logprobs:
      - token_id: 374      # " is"
        logprob: -0.18
      - token_id: 596      # "'s"
        logprob: -2.20
      - token_id: 13       # "."
        logprob: -3.50

The first dictionary supervises assistant_logits[0]; the second supervises assistant_logits[1]. For the first row, the helper exponentiates the teacher logprobs, renormalizes those three candidates into a small distribution over {12041, 987, 14321}, and penalizes the student if logits[0] does not put matching mass on those ids. The same operation is repeated for logits[1]. This is distillation rather than ordinary next-token cross-entropy: the student is not only told “the answer token was 12041”; it is also told that the teacher considered " London" much more plausible than " Lyon".

def _sparse_distillation_loss(
    logits: torch.Tensor,
    supervision: list[dict[str, Any]],
) -> torch.Tensor:
    log_probs = torch.log_softmax(logits.float(), dim=-1)
    token_losses: list[torch.Tensor] = []

    for row_idx, token_supervision in enumerate(supervision):
        candidate_ids: list[int] = []
        candidate_weights: list[float] = []

        for candidate in token_supervision["top_logprobs"]:
            token_id = candidate.get("token_id")
            if token_id is None:
                continue
            candidate_ids.append(int(token_id))
            candidate_weights.append(math.exp(float(candidate["logprob"])))

        target_token_id = int(token_supervision["token_id"])
        if target_token_id not in candidate_ids:
            candidate_ids.append(target_token_id)
            candidate_weights.append(math.exp(float(token_supervision["logprob"])))

        weights = torch.tensor(candidate_weights, device=logits.device, dtype=torch.float32)
        weights = weights / weights.sum()
        candidate_ids_tensor = torch.tensor(candidate_ids, device=logits.device)
        token_losses.append(
            -(weights * log_probs[row_idx, candidate_ids_tensor]).sum()
        )

    return torch.stack(token_losses).mean()

There are three small but important details in that helper. The sparse teacher mass comes from top_logprobs; the actual target token is inserted if it fell outside the top-k list; and the remaining candidates are renormalized before cross-entropy is computed.

The forward pass aligns the student logits with the assistant target tokens by teacher-forcing the answer prefix:

outputs = model(
    input_ids=model_input,
    past_key_values=cartridge.as_cache(model.config),
    use_cache=False,
)
target_len = len(example.assistant_token_ids)
start_idx = prompt_ids.shape[-1] - 1
end_idx = start_idx + target_len
assistant_logits = outputs.logits[0, start_idx:end_idx, :]
return _sparse_distillation_loss(assistant_logits, example.assistant_supervision)

The conceptual point is that the cartridge is trained through the frozen LLM. If moving one value vector helps the model put probability mass on the teacher’s answer token, gradient descent will move that vector. This is why a cartridge slot is not a token slot after training.

Cartridges Benchmarks From The Local Repo

The cartridges repo includes stable benchmark reports on small Wikipedia-backed QA tasks. The exact-match numbers are not the most important part here because the semantic judge is a better fit for these generated-answer runs. The systems numbers are the point: query-time prefill drops because the model no longer needs to prefill the full corpus on every question.

Cartridges are best read as an offline compilation strategy for a stable corpus. If the corpus is a codebase, legal document collection, textbook, or project memory that will receive many queries, a per-corpus optimization cost can be reasonable. If the corpus changes every request, the cost dominates.

STILL: Learn The Compressor Instead Of The Cache

Cartridges optimize $C_s$ directly for every source corpus $s$. STILL asks whether that optimization can be amortized.

Instead of solving:

for every new corpus, STILL learns a function:

Training the compactor $f_\phi$ is expensive, but it is reusable. For a new corpus, build the full cache once, run the compactor once, save the compact cache, and answer future questions against the compact artifact.

The Baseten write-up frames this as the missing amortization step for neural KV compaction. Cartridges validate that optimized compact caches can work, but they spend optimization on every new context. STILL spends optimization once on a compactor that can be reused across contexts.

In the local implementation, each transformer layer gets its own compactor:

self.layers = nn.ModuleList(
    [
        StillLayerCompactor(
            head_dim=head_dim,
            num_latents=num_latents,
            rope_theta=rope_theta,
        )
        for _ in range(num_hidden_layers)
    ]
)

The layer compactor takes:

and returns:

plus an additive attention bias:

That bias is important. The compact values carry content. The compact keys define where queries land. The bias lets the compactor adjust the prior preference over compact slots after compression.

STILL's compactor path: unrotate RoPE keys, let perceiver latents read the cache, project compact K/V, and add a β bias for compressed attention mass.

RoPE Makes Compaction Subtle

A raw key in a RoPE model is position-rotated. If we compress rotated keys directly, the compactor has to learn both content and position artifacts. STILL instead unrotates keys before feeding them to the perceiver, then rerotates compact keys at their new latent positions.

The implementation exposes this explicitly:

def apply_rope(
    x: torch.Tensor,
    positions: torch.Tensor,
    *,
    theta: float,
    inverse: bool = False,
) -> torch.Tensor:
    cos, sin = _rope_cos_sin(positions, dim=x.shape[-1], theta=theta)
    while cos.dim() < x.dim():
        cos = cos.unsqueeze(0)
        sin = sin.unsqueeze(0)
    if inverse:
        sin = -sin
    x_float = x.to(torch.float32)
    return ((x_float * cos) + (_rotate_half(x_float) * sin)).to(x.dtype)

Inside StillLayerCompactor.forward, latent_positions is not just metadata. It is threaded through the perceiver blocks and then used again when the compact keys are rotated back into the coordinate system expected by the frozen LLM:

token_positions = torch.arange(seq_len, device=keys.device, dtype=torch.long)
latent_positions = self._latent_positions(seq_len, keys.device)

unrotated_keys = apply_rope(keys, token_positions, theta=self.rope_theta, inverse=True)
kv_input = torch.cat([unrotated_keys, values], dim=-1)
kv_input = kv_input.squeeze(0).reshape(num_heads, seq_len, head_dim * 2).to(module_dtype)
latents = self.latents.unsqueeze(0).expand(num_heads, -1, -1).to(module_dtype)

for block in self.blocks:
    latents = block(
        latents,
        kv_input,
        latent_positions=latent_positions,
        token_positions=token_positions,
    )

compact_keys = self.key_head(latents)
compact_values = self.value_head(latents)
compact_biases = self.bias_head(latents).squeeze(-1)
compact_keys = apply_rope(compact_keys, latent_positions, theta=self.rope_theta)

The latent positions are evenly spaced across the original sequence:

def _latent_positions(self, seq_len: int, device: torch.device) -> torch.Tensor:
    if self.num_latents == 1:
        return torch.zeros(1, device=device, dtype=torch.long)
    values = torch.linspace(
        0,
        max(seq_len - 1, 0),
        steps=self.num_latents,
        device=device,
        dtype=torch.float32,
    )
    return values.round().to(torch.long)

This is a simple way to say: compact slot $j$ represents some position in the original timeline. It is not the only possible scheme, but it gives the frozen model position-compatible keys.

STILL’s Perceiver Block

The compactor uses learned latents. For each KV head, those latents cross-attend into the dense cache, then self-attend with one another.

One block does:

where:

The source code mirrors this:

latents = self.cross_norm(latents + cross_out)
latents = self.self_norm(latents + self.self_attn(latents))

The cross-attention uses learned latent queries and the dense cache as memory:

q = apply_rope(self.q_proj(latents), latent_positions, theta=self.rope_theta)
k = apply_rope(self.k_proj(kv_input), token_positions, theta=self.rope_theta)
v = self.v_proj(kv_input)
scale = 1.0 / math.sqrt(self.dim)
weights = torch.softmax(torch.matmul(q, k.transpose(-1, -2)) * scale, dim=-1)
outputs = self.out_proj(torch.matmul(weights, v))

After two perceiver blocks, the output heads produce compact keys, values, and biases:

compact_keys = self.key_head(latents)
compact_values = self.value_head(latents)
compact_biases = self.bias_head(latents).squeeze(-1)
compact_keys = apply_rope(compact_keys, latent_positions, theta=self.rope_theta)

The initialization is deliberately structured. The key head initially reads the first half of the latent vector, the value head reads the second half, and the bias head starts at zero:

self.key_head.weight.data.zero_()
self.key_head.weight.data[:, : self.head_dim] = torch.eye(self.head_dim)
self.value_head.weight.data.zero_()
self.value_head.weight.data[:, self.head_dim :] = torch.eye(self.head_dim)

That makes early behavior easier to reason about: the compactor starts closer to a copy-style mapping than arbitrary noise.

STILL Training Objective

The teacher is the frozen model with the full cache. The student is the same frozen model with the compact cache produced by the compactor. Training updates only the compactor.

The general objective in the repo is:

The implementation supports both terms:

def _distillation_loss(
    *,
    teacher_logits: torch.Tensor,
    student_logits: torch.Tensor,
    target_token_ids: list[int],
    kl_weight: float,
    exact_token_ce_weight: float,
) -> torch.Tensor:
    loss = student_logits.new_tensor(0.0, dtype=torch.float32)
    if kl_weight > 0.0:
        loss = loss + (kl_weight * _kl_loss(teacher_logits, student_logits))
    if exact_token_ce_weight > 0.0:
        ce_term = _exact_token_ce_loss(student_logits, target_token_ids)
        loss = loss + (exact_token_ce_weight * ce_term)
    return loss

The teacher and student paths are computed in the same function:

with torch.no_grad():
    full_outputs = model(input_ids=context_ids, use_cache=True)
    teacher_outputs = model(
        input_ids=model_input,
        past_key_values=full_outputs.past_key_values,
        use_cache=False,
    )

compact_cache = compactor(full_outputs.past_key_values)
student_outputs = model(
    input_ids=model_input,
    past_key_values=compact_cache.as_cache(model.config),
    still_layer_biases=compact_cache.biases,
    use_cache=False,
)

This is the crucial amortization step. During training, the compactor sees many contexts and learns how to compress caches. During deployment, a new context only needs:

with torch.no_grad():
    outputs = model(input_ids=context_ids, use_cache=True)
    cache = compactor(outputs.past_key_values)

That is one full prefill plus one compactor pass, not hundreds of gradient steps on a new cache.

STILL Benchmark From The Local Repo

The STILL repo includes a final MCQ benchmark using Qwen/Qwen3-4B, 1024 compact latents, 115 training Wikipedia articles, and 20 held-out Wikipedia articles with 10 MCQ questions per held-out article.

The benchmark compares full context, truncation, STILL, and a cartridge baseline under aligned metric definitions.

The result is not “STILL beats cartridges on quality.” It does not in this benchmark. Cartridges is much stronger on held-out MCQ accuracy because it optimizes a compact cache for each target page.

The result is: STILL shows the systems shape of a reusable compactor. The online latency is lowest, the per-target build is small, and the one-time training cost is separated from per-target preparation. The quality gap is the central open engineering problem in this reproduction.

Cartridges Versus STILL

The easiest way to compare them is by where the optimization lives.

Mathematically:

Cartridges are like per-document prompt tuning, except the prompt lives in KV space. STILL is like learning an encoder that produces those prompts directly.

Operationally:

What Changes In Practice

The last piece is to translate the taxonomy back into engineering tradeoffs: what is different from text prompt compression, what the implementations taught, and where the next measurements should go.

Why This Is Not Just Prompt Compression

Prompt compression methods usually emit text. The model then tokenizes that text and builds a normal KV cache. That is useful, but it puts the bottleneck through language. If a detail is not in the summary, it is gone.

KV compaction emits model-internal vectors. This gives it a different kind of capacity. A compact slot can store information in a way that is not a grammatical sentence. It can act as a router, a value carrier, or an attention prior.

This is also why the methods are harder to debug. When a textual summary fails, we can read it. When a compact value vector fails, we need probes: attention maps, nearest-neighbor analyses, teacher-student token deltas, ablations by layer, and task-specific recall tests.

Implementation Lessons

The code across these repos makes a few practical lessons clear.

1. KV accounting should be canonical. Use the tensor formula, not allocator memory, for method comparisons:

Allocator memory still matters for deployment, but canonical bytes make algorithmic compression ratios comparable.

2. Attention-score methods need score visibility. SnapKV and H2O need attention matrices or attention history. That can conflict with kernels designed to avoid materializing attention probabilities.

3. Positional encoding is part of the cache. STILL’s unrotate/compress/rerotate path exists because a key vector is not merely content. It is content after positional rotation.

4. “Compact tokens” do not imply equal serving cost. In the STILL benchmark, truncation and STILL both operate around a 1024-token budget, but truncation still reprefills the retained prompt on every query. STILL pays a target cache build once, then uses a short query prompt with the compact cache already loaded.

5. Decode length can distort latency. The STILL report explicitly tracks generated token counts because methods can normalize to the same MCQ answer while producing different raw completions.

A Minimal Pseudocode Summary

Attention Sink:

prefill(K, V, T):
    if T <= G + W:
        keep all
    else:
        keep [0:G] and [T-W:T]

decode(k_t, v_t):
    keep [0:G]
    replace oldest / lowest-position non-global slot with (k_t, v_t)

L2 norm pruning:

score_i = -||k_i||_2
keep top B positions by score
append future decode tokens into freed slots

SnapKV:

run prompt prefill
observe attention from last W prompt tokens to earlier prompt tokens
score_i = pooled_sum_attention_to_i
keep top (B - W) scored older tokens plus the last W tokens

H2O:

for each decode step:
    update average attention history for cached tokens
    if cache is full:
        evict token with lowest history score
    insert current token

Cartridges:

for each corpus:
    C = KV cache from first p prefix tokens
    freeze base model
    repeat gradient steps:
        teacher = model(full_context, query, answer_prefix)
        student = model(query, answer_prefix, past_key_values=C)
        update C to match teacher answer distribution
    save C

STILL:

train once:
    for each training corpus:
        full_cache = frozen_model.prefill(corpus)
        compact_cache = compactor(full_cache)
        teacher = frozen_model(query, past_key_values=full_cache)
        student = frozen_model(query, past_key_values=compact_cache)
        update compactor to match teacher / target answer

serve new corpus:
    full_cache = frozen_model.prefill(corpus)
    compact_cache = compactor(full_cache)
    answer future queries from compact_cache

Where I Would Push Next

The current implementations make the core ideas inspectable. The natural next steps are less about adding another eviction heuristic and more about measurement.

For Cartridges, I would want:

• layer-wise ablations: which layers need learned values versus mostly learned keys?
• query-count break-even curves under realistic serving assumptions;
• composition tests where multiple cartridges are loaded together;
• probes that distinguish factual storage from routing behavior.

For STILL, I would want:

• stronger training distributions than small MCQ supervision;
• KL-heavy or continuation-token training once decode collapse is controlled;
• different latent-position schedules;
• ablations for beta, RoPE unrotation, and identity initialization;
• iterative compaction tests where compact memory is prepended to the next chunk and compacted again.

For token eviction, I would want:

• apples-to-apples results under the same attention kernel constraints;
• memory accounting that separates effective cache size from allocated buffer size;
• long-horizon stability tests, not just short prompt generation.

Takeaway

KV cache compression is not one trick. There are at least three levels:

1. Eviction: keep a smaller subset of original tokens.
2. Per-corpus compaction: optimize compact K/V tensors for one source.
3. Amortized compaction: train a reusable compressor that builds compact K/V tensors for new sources.

Attention Sink, L2 pruning, SnapKV, and H2O are useful because they make bounded-cache inference possible with relatively small changes. Cartridges are useful because they show that a compact KV object can behave like a corpus-specific memory. STILL is useful because it points toward the systems shape we actually want: compress a new context in one forward-pass-like build step, then reuse that compact memory across future queries.

The hard part is preserving the full-context model’s behavior while paying less than full-context serving cost. That is exactly why this area is interesting: it sits at the boundary between representation learning, attention mechanics, and inference systems.

Follow me on Twitter, Github or connect on LinkedIn.

Paper: Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models
Official implementation: DeepSeek-AI/Engram

Attention is not memory

Self-attention can resolve relationships in a sentence. It does not automatically provide a grounded representation of what the entities actually are.

In a standard Transformer, this grounding is reconstructed through repeated computation. Attention composes nearby tokens. Feed-forward layers transform features. Later layers gradually turn surface strings into semantic representations.

The Engram paper frames this as an architectural mismatch: dynamic reasoning should use computation, while common static phrases should often use lookup.

The FFN already looks like a memory

A Transformer MLP can be read as a bank of pattern detectors and value writers.

$$\operatorname{FFN}(h) = W_{\text{down}} \, \sigma(W_{\text{up}}h + b_{\text{up}}) + b_{\text{down}}.$$

Geva et al. showed that FFNs behave like key-value memories: rows of $W_{\text{up}}$ detect patterns, while columns of $W_{\text{down}}$ write value vectors into the residual stream.

MoE scales this by adding many FFNs and routing each token to a few experts:

$$\operatorname{MoE}(h_t)=\sum_{i \in \operatorname{TopK}(r(h_t))} p_i(h_t)E_i(h_t).$$

Static facts want tables

For a single token, lookup is simple. For phrases, the combinatorics explode.

A token embedding table maps an ID directly to a vector:

$$e = E[x], \qquad E \in \mathbb{R}^{|V| \times d}.$$

But facts are usually phrase-level. "Harry" is ambiguous; "Harry Potter" is a much more specific key.

Harry Potter

Φ

A direct bigram table with $|V|=128{,}000$ would have:

$$|V|^2 = 128{,}000^2 = 16{,}384{,}000{,}000.$$

Hash the local phrase

Engram compresses token IDs, hashes suffix n-grams, and retrieves rows from multiple embedding tables.

First, a tokenizer projection maps raw token IDs into canonical IDs:

$$P: V \to V', \qquad x'_t = P(x_t).$$

Then Engram forms suffix n-grams:

$$g_{t,n} = (x'_{t-n+1}, \ldots, x'_t).$$

Each hash head maps the compressed n-gram into a table row:

$$z_{t,n,k} = \phi_{n,k}(g_{t,n}), \qquad e_{t,n,k} = E_{n,k}[z_{t,n,k}].$$ $$e_t=\big\Vert_{n=2}^{N}\big\Vert_{k=1}^{K}e_{t,n,k}.$$

Why multiplicative-XOR?

Addition creates structured collisions and loses order. Plain XOR also loses order because it is commutative. Engram uses position-specific multipliers before XOR:

$$\phi_{n,k}(g_{t,n})= \left(\bigoplus_{i=0}^{n-1}m^{(\ell,k)}_i x'_{t-i}\right)\bmod M_{n,k}.$$

Lookup needs a gate

Static memory is useful only when the current context agrees with it.

The retrieved vector $e_t$ is projected into a key and value:

$$k_t = W_K e_t, \qquad v_t = W_V e_t.$$

The hidden state is the query. The scalar gate is:

$$\alpha_t=\sigma\left( \frac{\operatorname{RMSNorm}(h_t)^\top\operatorname{RMSNorm}(k_t)}{\sqrt{d}} \right).$$ $$\tilde{v}_t=\alpha_t v_t.$$

Short convolution

After gating, Engram applies a short depthwise causal convolution and a residual path:

$$Y=\operatorname{SiLU}\left(\operatorname{Conv1D}(\operatorname{RMSNorm}(\tilde{V}))\right)+\tilde{V}.$$ $$H^{(\ell)} \leftarrow H^{(\ell)} + Y.$$

Inside the Transformer, not just at the input

Engram is inserted into selected Transformer blocks. The paper's 27B model uses layers 2 and 15.

Layer 1 is too raw

The hidden state is still close to token embeddings, so context-aware gating has little context to use.

Layer 2 is the sweet spot

One round of attention is enough to make the gate useful while still being early enough to save depth.

Middle layers refine

A later Engram module catches associations that only become clear after partial processing.

For multi-branch mHC backbones, Engram shares the memory table and value projection, but uses branch-specific key projections:

$$\alpha^{(m)}_t= \sigma\left( \frac{\operatorname{RMSNorm}(h^{(m)}_t)^\top\operatorname{RMSNorm}(W^{(m)}_K e_t)}{\sqrt{d}} \right),$$ $$u^{(m)}_t=\alpha^{(m)}_t(W_Ve_t).$$

How much memory is enough?

Engram's strongest empirical claim is that sparse capacity should be split between MoE and memory.

$$P_{\text{sparse}}=P_{\text{tot}}-P_{\text{act}}.$$ $$P_{\text{MoE}}^{(\text{sparse})}=\rho P_{\text{sparse}}, \qquad P_{\text{Engram}}=(1-\rho)P_{\text{sparse}}.$$

Sparsity allocation

Move the slider. The paper's optimum appears around $\rho \approx 0.75$ to $0.80$, where most sparse capacity remains MoE but a meaningful chunk becomes Engram memory.

rho = 0.80

MoE

Engram

MoE sparse share 80%

Engram sparse share 20%

Toy validation loss 1.711

What changes at scale?

Engram-27B is iso-parameter and iso-FLOPs relative to MoE-27B. The win comes from reallocating sparse capacity, not from spending more activated compute.

Model	Total params	Activated params	Experts	Engram params
Dense-4B	4.1B	3.8B	none	none
MoE-27B	26.7B	3.8B	2 shared + 72 routed, top-6	none
Engram-27B	26.7B	3.8B	2 shared + 55 routed, top-6	5.7B
Engram-40B	39.5B	3.8B	2 shared + 55 routed, top-6	18.5B

Effective depth

Engram helps shallow layers behave like deeper MoE layers because static local reconstruction is handled by lookup.

$$ \text{lookup for static facts} \Rightarrow \text{less early reconstruction} \Rightarrow \text{more effective depth} $$

class TransformerBlock(nn.Module):
    def forward(self, input_ids, hidden_states):
        if self.engram is not None:
            hidden_states = (
                self.engram(hidden_states=hidden_states, input_ids=input_ids)
                + hidden_states
            )

        hidden_states = self.attn(hidden_states) + hidden_states
        hidden_states = self.moe(hidden_states) + hidden_states
        return hidden_states

old2new = {}
key2new = {}

for tid in range(vocab_size):
    text = tokenizer.decode([tid], skip_special_tokens=False)
    key = token_string_if_undecodable(text) or normalize(text)

    if key not in key2new:
        key2new[key] = len(key2new)

    old2new[tid] = key2new[key]

lookup = np.empty(vocab_size, dtype=np.int64)
for tid in range(vocab_size):
    lookup[tid] = old2new[tid]

def hash(self, input_ids):
    input_ids = self.compressed_tokenizer(input_ids)
    hash_ids_for_all_layers = {}

    for layer_id in self.layer_ids:
        hash_ids_for_all_layers[layer_id] = self._get_ngram_hashes(
            input_ids,
            layer_id=layer_id,
        )

    return hash_ids_for_all_layers

def shift_k(k):
    if k == 0:
        return x
    shifted = np.pad(
        x,
        ((0, 0), (k, 0)),
        mode="constant",
        constant_values=self.pad_id,
    )[:, :T]
    return shifted

base_shifts = [shift_k(k) for k in range(self.max_ngram_size)]

for n in range(2, self.max_ngram_size + 1):
    n_gram_index = n - 2
    tokens = base_shifts[:n]

    mix = tokens[0] * multipliers[0]
    for k in range(1, n):
        mix = np.bitwise_xor(mix, tokens[k] * multipliers[k])

    for j, mod in enumerate(head_vocab_sizes):
        head_hash = mix % int(mod)
        all_hashes.append(head_hash.astype(np.int64, copy=False))

return np.stack(all_hashes, axis=2)

offsets = [0]
for table_size in list_of_N[:-1]:
    offsets.append(offsets[-1] + table_size)

shifted_input_ids = input_ids + self.offsets
rows = self.embedding(shifted_input_ids)

hash_input_ids = torch.from_numpy(
    self.hash_mapping.hash(input_ids)[self.layer_id]
)

embeddings = self.multi_head_embedding(hash_input_ids)
embeddings = embeddings.flatten(start_dim=-2)

gates = []
for hc_idx in range(backbone_config.hc_mult):
    key = self.key_projs[hc_idx](embeddings)
    normed_key = self.norm1[hc_idx](key)

    query = hidden_states[:, :, hc_idx, :]
    normed_query = self.norm2[hc_idx](query)

    gate = (normed_key * normed_query).sum(dim=-1)
    gate = gate / math.sqrt(backbone_config.hidden_size)
    gate = gate.abs().clamp_min(1e-6).sqrt() * gate.sign()
    gate = gate.sigmoid().unsqueeze(-1)
    gates.append(gate)

gates = torch.stack(gates, dim=2)
value = gates * self.value_proj(embeddings).unsqueeze(2)

output = value + self.short_conv(value)
return output

x

Follow me on Twitter, Github or connect on LinkedIn.

Paper: Physics of Language Models: Part 4.1, Architecture Design and the Magic of Canon Layers
Official implementation: facebookresearch/PhysicsLM4

Canon Layers are a small architectural primitive from Zeyuan Allen-Zhu’s Physics of Language Models: Part 4.1. The basic idea is simple: give every token a cheap causal path to nearby past token states.

That path is not meant to replace attention. It handles a different job.

Attention should spend capacity on content-addressed routing and retrieval. It should not have to spend layers on routine neighbor-to-neighbor transport.

The mechanism is a residual causal depthwise convolution over the sequence axis. It is local, cheap, and easy to insert into existing Transformer, linear-attention, and SSM-style blocks.

The missing path in a standard Transformer

A pre-norm Transformer block usually has:

This gives two strong paths:

• a vertical residual path, where token position $t$ preserves and refines its own representation across layers;
• a global attention path, where token position $t$ can retrieve content from previous positions.

But the MLP is pointwise over tokens. It mixes channels, not positions. Attention can move information from $t-1$ to $t$, but attention is a global content-routing mechanism. Using it for routine local relay is expensive and depth-inefficient.

Canon adds a third path:

That is why the paper describes Canon as horizontal information flow. The ordinary residual stream is vertical across depth; Canon is local residual flow across positions.

Associative recall shows the problem

Consider the causal sequence:

[A] [B] ... [A] [?]

The desired next token is [B]. A natural mechanism is:

1. the second [A] attends to the first [A];
2. the representation at the first [A] carries enough information to identify the following [B];
3. the model predicts [B].

The catch is causal masking. The first [A] cannot see its future neighbor [B] at the same layer. A model often needs one operation to move information locally from [B] into a neighboring representation, then another operation to retrieve it globally.

Canon makes that first local enrichment cheap.

The Canon operator

Let a sequence of hidden states be:

A width-4 Canon layer computes:

where $w_r\in\mathbb{R}^{m}$ are learned channelwise weights, $\odot$ is elementwise multiplication, and missing past states are zero-padded.

The residual form is:

Equivalently, for batch index $b$, position $t$, channel $c$, and kernel size $K$:

with $x_{b,t-r,c}=0$ when $t-r<0$.

The key word is depthwise. Channel $c$ reads only channel $c$ over nearby positions. Canon does not perform hidden-dimension mixing; the projections and MLP still own that job.

Why the residual matters

Without the residual path:

With the residual path:

The residual version is easier to insert because it starts as a local perturbation around the existing representation. If the local signal is useful, the model can add it. If it is not useful, the model can learn small weights without destroying the vertical residual stream.

The Canon paper’s ablations report that residual Canon is materially more stable and efficient than non-residual variants. The implementation also exposes canon_residual as a configuration flag, with the released LlamaCanon path defaulting to residual behavior.

Canon is not local attention

Local attention computes content-dependent weights:

Canon computes fixed learned local propagation:

The distinction matters:

Canon is closer to a short learned transport operator than to a retrieval mechanism.

Where Canon goes in a Transformer block

The paper studies four insertion points. For hidden width $d$, Canon-ABCD means:

Canon-A: before attention

Then:

Attention receives token states that already contain a short causal neighborhood.

Canon-B: inside attention

After Q/K/V projection:

Canon-B applies local mixing to that projected representation:

For ordinary MHA with equal Q/K/V widths, $m=3d$. For grouped-query attention:

The released LlamaCanon code computes exactly this total dimension before constructing canonB.

Canon-C: before the MLP

The MLP receives a locally enriched representation.

Canon-D: inside the MLP

For a gated MLP:

LlamaCanon concatenates the gate and up projections:

then computes:

For a Llama-style gated MLP with intermediate width $\frac{8}{3}d$, Canon-D has width:

Canon-ABCD pseudocode

The same residual local mixer appears at different internal representations:

def canon_residual(x, canon_conv):
    # x: [batch, seq, channels]
    # canon_conv: causal depthwise Conv1d over the sequence dimension
    return x + canon_conv(x)

For a Llama-style pre-norm block with grouped-query attention and a gated MLP:

def llama_block_with_canon(x, mask=None, cache=None):
    # x: [B, T, d]

    residual = x
    h = rmsnorm_attn(x)

    if canonA is not None:
        h = canon_residual(h, canonA)          # [B, T, d]

    q = q_proj(h)
    k = k_proj(h)
    v = v_proj(h)

    if canonB is not None:
        qkv = concat([q, k, v], dim=-1)
        qkv = canon_residual(qkv, canonB)
        q, k, v = split(qkv, [q_dim, k_dim, v_dim], dim=-1)

    q = apply_rope(q)
    k = apply_rope(k)
    a = causal_attention(q, k, v, mask=mask, cache=cache)
    x = residual + o_proj(a)

    residual = x
    h = rmsnorm_mlp(x)

    if canonC is not None:
        h = canon_residual(h, canonC)          # [B, T, d]

    gate = gate_proj(h)
    up = up_proj(h)

    if canonD is not None:
        z = concat([gate, up], dim=-1)
        z = canon_residual(z, canonD)
        gate, up = z.chunk(2, dim=-1)

    x = residual + down_proj(silu(gate) * up)
    return x

Partial variants such as Canon-AC, Canon-ACD, or Canon-ABC are also meaningful. The paper’s ablations find that the benefits are cumulative, and that Canon-ACD can help even without modifying the attention projections.

Tensor shapes for the core mixer

The minimal PyTorch version for a [B,T,D] tensor is:

import torch
import torch.nn as nn
import torch.nn.functional as F


class CanonResidualMixer(nn.Module):
    def __init__(self, channels: int, kernel_size: int = 4):
        super().__init__()
        self.kernel_size = kernel_size
        self.conv = nn.Conv1d(
            in_channels=channels,
            out_channels=channels,
            kernel_size=kernel_size,
            groups=channels,
            bias=False,
        )

    def forward(self, x):
        # x: [B, T, D]
        xt = x.transpose(1, 2)                    # [B, D, T]
        xt = F.pad(xt, (self.kernel_size - 1, 0)) # [B, D, T + K - 1]
        mixed = self.conv(xt)                     # [B, D, T]
        mixed = mixed.transpose(1, 2)             # [B, T, D]
        return x + mixed

Shape summary:

For $K=4$, one channel computes:

Then:

Why groups=channels matters

With depthwise convolution:

nn.Conv1d(D, D, K, groups=D)

the parameter tensor has shape:

so parameters scale as:

With full convolution:

nn.Conv1d(D, D, K, groups=1)

the parameter tensor has shape:

so parameters scale as:

For $D=4096$ and $K=4$:

That gap is the reason Canon isolates local sequence transport from channel mixing. Channel mixing remains in the projections and MLP, where it already exists.

Complexity and runtime

For batch $B$, sequence length $T$, hidden width $D$, and small kernel $K$:

The attention matrix/value aggregation term is roughly:

plus projection costs.

Asymptotically, Canon is tiny. Practically, it is not free: every additional operator can add memory movement and kernel-launch overhead. The Part 4.1 paper reports that Canon-ABCD adds fewer than $0.45\%$ parameters for GPT-2-small, and for a 1.3B Llama-style model it adds about $0.0063\%$ parameters. The same footnote reports nonzero naive H100 runtime overheads, with Canon-AC cheaper than Canon-ABCD.

The released code uses a ShortConvolution wrapper with causal_conv1d when available and when the kernel is in ${2,3,4}$. During generation, the convolution cache stores only the last $K$ states per channel:

The synthetic playground

The paper argues that academic-scale real-data pretraining can be too noisy for architecture science. Perplexity mixes many skills together; benchmark swings can hide whether an architecture improved reasoning, knowledge storage, local composition, or something else.

The Part 4.1 experiments therefore use five controlled synthetic pretraining tasks:

The point is not that synthetic tasks are the final benchmark. They isolate mechanisms. If a change improves Depo but not Capo, or helps NoPE but not RoPE, the result is easier to interpret than a single mixed-corpus loss number.

Depo: depth

Depo builds a directed permutation from key-value pairs:

<bos> x1 y1 x2 y2 ... xn yn <query_k> q <ans> a <eos>

If the pairs define $f(x_i)=y_i$, the target is:

The model must compute the $k$-hop successor internally, without writing intermediate chain-of-thought tokens. Depo2 makes each node span multiple tokens, so a 4-token Canon window cannot solve the task by direct copying. The local mixer must improve segment representations that attention can later chain globally.

Brevo: breadth

Brevo gives the model a directed acyclic graph and asks for recursive dependencies in topological order. The hard part is not one long chain; it is parallel dependency processing across branches.

Capo: capacity

Capo measures reliable storage of synthetic facts, often as bits per parameter. Limited-exposure regimes are important because overtraining can hide architectural differences.

Mano: manipulation

Mano uses modular arithmetic expressions. The model must retrieve learned operation tables and compose them internally. This tests manipulation of knowledge stored in weights rather than only information present in the prompt.

Lano: structure

Lano uses CFG-like sequences with local ambiguity. Correct prediction can require maintaining recursive global structure rather than memorizing nearby tokens.

What the results imply

For Transformer-style models, the Part 4.1 paper reports that Canon-ABCD improves reasoning depth by roughly $2$-$4\times$ in the controlled setup, reasoning breadth by about $30\%$, knowledge manipulation length by about $30\%$, and knowledge capacity in limited-exposure factual-storage regimes.

The strongest interpretation is not that Canon solves every task inside a four-token window. It is that better local representations make later global routing easier.

The NoPE result is especially interesting. NoPE means no positional embedding. Without positional encoding, a Transformer has weak order information. With Canon, NoPE becomes far stronger, often competitive with RoPE+Canon in the reported synthetic setup. A causal convolution injects order-sensitive local structure:

The paper also studies partial RoPE. With Canon present, reduced-RoPE variants can work well, which matters because heavy RoPE usage can hurt length generalization.

Linear models and SSMs

The paper compares Transformers, GLA, Mamba2, and GDN under the same synthetic tasks. A useful takeaway is that local convolution-like components inside some linear/SSM architectures already explain a lot of their behavior.

In the paper’s terminology:

• Mamba2’s internal conv1d resembles a partial non-residual Canon-B;
• GLA and GDN implementations also contain conv-like local components;
• adding Canon systematically makes comparisons fairer because every model receives the same local-transport primitive.

After adding Canon broadly, linear models still tend to lag full-attention Transformers on deep retrieval-heavy reasoning. The diagnosis is not only state size. The harder problem is memory dynamics: compressed recurrent state must preserve and retrieve fine-grained facts across multiple hops without compounding errors.

The Part 4.2 code release extends the story to real-world pretraining recipes and released model families, including LlamaCanon, GLA, GDN, and Mamba2 variants.

Canon versus related mechanisms

Primer

Primer introduced squared ReLU and a depthwise convolution after Q/K/V projection. The Q/K/V convolution part is closest to Canon-B without the residual path:

Canon generalizes the idea in three ways:

1. it adds an explicit residual around the local mixer;
2. it applies the primitive at A/B/C/D, not only Q/K/V;
3. it studies the primitive across Transformers, linear attention, and SSM-style models.

Longformer-style local attention

Longformer sparsifies attention with sliding windows and task-specific global attention. Canon works on a different axis.

Local attention asks:

Canon asks:

They can coexist.

Mamba2

Mamba2 is built around state-space duality and selective SSM computation. Its local convolution is a frontend to a recurrent/SSM memory system:

Canon isolates the local convolutional part as a reusable residual primitive that can be applied outside a specific SSM block.

Uniform attention

Earlier Physics of Language Models work found that uniform averaging over recent tokens could help CFG-style tasks. Canon can be viewed as a learned, channelwise, modular version of that local averaging:

Implementation details from LlamaCanon

The released LlamaCanon helper uses a ShortConvolution module:

ShortConvolution(
    hidden_size=dim,
    kernel_size=config.canon_kernel,
    bias=config.canon_bias,
    activation="silu" if config.canon_activation else None,
    use_fast_conv1d=causal_conv1d_available and config.canon_kernel in [2, 3, 4],
)

It is dimension-last at the interface:

then rearranges internally to Conv1d layout:

The helper masks padded positions, uses the fast causal_conv1d kernel when available, and supports decode-time cache updates through a [B,D,K] state.

The code exposes:

• canon_set, selecting any subset of A, B, C, D;
• canon_kernel, usually $4$;
• canon_residual, controlling whether the output is hidden_states + hidden_states2;
• canon_activation, available but not recommended by the paper for Transformer Canon layers;
• canon_bias, generally avoided.

For packed or padded batches, Canon must respect the same valid-token mask as attention. Otherwise, a causal convolution can propagate padding artifacts into valid positions.

Practical choices

Initialization

There are several reasonable options:

1. Default initialization. This matches the released implementation path.

2. Zero initialization. This makes Canon an exact identity at step zero:

   \[y=x+0=x.\]

   That is useful when retrofitting Canon into an already trained model.

3. Past-average initialization. For $K=4$, initialize previous offsets to $\frac13$ and current offset to $0$:

   \[y_t=x_t+\frac13(x_{t-1}+x_{t-2}+x_{t-3}).\]

   This tests the local-context hypothesis directly, but it is a design choice rather than the default released setup.

Causal padding

Use left padding:

F.pad(x, (K - 1, 0))

Right padding would either shift outputs incorrectly or leak future information.

Optimized kernels

torch.nn.Conv1d is the generic API. It is not automatically the optimized Dao-AILab causal-conv1d path. The implementation must call that package explicitly, as LlamaCanon’s helper does.

Open engineering questions

Runtime overhead

Parameter overhead is tiny, but runtime overhead is still real. Multiple small convolutions can add memory traffic and kernel launches. A production implementation would likely fuse Canon with adjacent projections or batch several Canon calls together.

Dynamic Canon

The studied operator uses fixed learned weights. A dynamic version could use input-conditioned local weights:

That moves Canon closer to lightweight local attention. It may improve expressivity, but it also changes the clean cost and interpretation.

MoE interaction

Canon-D inside a mixture-of-experts MLP is awkward because neighboring tokens may be routed to different experts. Canon-ABC is easier; Canon-D requires a more careful dispatch design.

Long-range compression

Canon improves local flow. It does not remove the hard problem of preserving high-fidelity information through compressed recurrent state or across very long contexts. The paper’s linear-model results still suggest that full attention remains stronger for some deep in-context reasoning tasks.

Summary

Canon Layers are lightweight residual causal convolutions over neighboring token representations:

The architecture split is clean:

• attention handles content-addressed global routing;
• MLPs handle channelwise nonlinear transformation;
• Canon handles cheap local token-to-token propagation.

The empirical claim from the Canon paper is that this small primitive improves controlled measures of reasoning depth, reasoning breadth, knowledge manipulation, NoPE viability, and several linear/SSM architectures. The implementation claim is equally simple: Canon is depthwise causal Conv1D with a residual path, placed at selected A/B/C/D points inside a block.

Canon is not interesting because “convolution is back.” It is interesting because local horizontal flow is useful enough to deserve its own architectural slot.

References

• Zeyuan Allen-Zhu, Physics of Language Models: Part 4.1, Architecture Design and the Magic of Canon Layers, 2025.
• Zeyuan Allen-Zhu, Physics of Language Models: Part 4.2, Canon Layers at Scale where Synthetic Pretraining Resonates in Reality, 2025.
• facebookresearch, PhysicsLM4 code release.
• David R. So et al., Primer: Searching for Efficient Transformers for Language Modeling, NeurIPS 2021.
• Tri Dao and Albert Gu, Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality, ICML 2024.
• Iz Beltagy, Matthew E. Peters, and Arman Cohan, Longformer: The Long-Document Transformer, 2020.

Follow me on Twitter, Github or connect on LinkedIn.

Papers: Hyper-Connections and mHC: Manifold-Constrained Hyper-Connections

Hyper-Connections (HC) generalize the residual stream into multiple learned streams. Manifold-Constrained Hyper-Connections (mHC) keep that extra routing capacity while constraining the residual mixing matrices so deep products remain stable.

The baseline: what residual connections are really buying us

A standard residual layer is:

where:

• $\mathbf{x}_l \in \mathbb{R}^{C}$ is the residual-stream state at layer $l$,
• $\mathcal{F}$ is the layer body, e.g. attention or MLP,
• $\mathcal{W}_l$ are the layer weights.

The important part is not just the addition. It is the identity path. If we recursively expand from a shallow layer $l$ to a deeper layer $L$, we get:

The shallow signal $\mathbf{x}_l$ reaches layer $L$ unchanged. This gives deep networks a stable highway for forward activations and backward gradients.

Why this matters:

• without a skip path, deep networks must repeatedly multiply through learned transformations, so signals can vanish or explode;
• with the identity path, the model can start near a shallow network and learn residual corrections;
• in Transformers, the residual stream is the persistent memory channel through which attention and MLP blocks communicate.

The classical residual connection is stable, but rigid. Every layer receives exactly one stream and writes back to exactly one stream.

The problem HC was trying to solve

Residual variants such as Pre-Norm and Post-Norm make a trade-off:

• Pre-Norm stabilizes gradients, but can make adjacent deep representations overly similar. This is the “representation collapse” side of the trade-off.
• Post-Norm can improve representational separation, but is more prone to gradient vanishing.

The HC paper frames this as a “seesaw”: fixed residual wiring chooses one point on the stability-vs-expressivity curve. Hyper-Connections ask: can the network learn how strongly layers should connect, instead of hard-coding the residual topology?

The key HC move is to widen the residual stream.

Instead of carrying one vector:

HC carries $n$ parallel residual streams:

Here $n$ is the expansion rate. In the DeepSeek mHC paper, experiments commonly use $n=4$.

The layer body $\mathcal{F}$ still expects a normal $C$-dimensional input, so HC needs three maps:

1. Pre / aggregation map $\mathcal{H}_l^{\mathrm{pre}} \in \mathbb{R}^{1 \times n}$: combines $n$ streams into one layer input.
2. Post / expansion map $\mathcal{H}_l^{\mathrm{post}} \in \mathbb{R}^{1 \times n}$: writes the layer output back into $n$ streams.
3. Residual mixing map $\mathcal{H}_l^{\mathrm{res}} \in \mathbb{R}^{n \times n}$: mixes the residual streams directly.

The single-layer HC update is:

This is the same residual idea, but generalized:

• $\mathcal{H}_l^{\mathrm{pre}}\mathbf{x}_l$ chooses what mixture of streams enters attention or MLP;
• $\mathcal{H}_l^{\mathrm{post}\,\top}\mathcal{F}(\cdot)$ chooses where the new layer output is written;
• $\mathcal{H}_l^{\mathrm{res}}\mathbf{x}_l$ lets old streams exchange information before the new residual update is added.

In implementation terms, a single HC residual layer looks like this:

def hyper_connection_layer(x, layer_fn, h_pre, h_post, h_res):
    """
    x:      [n, C] multi-stream residual state
    h_pre:  [n]    aggregates streams into one layer input
    h_post: [n]    writes layer output back to streams
    h_res:  [n, n] mixes streams directly
    """
    layer_input = h_pre @ x                 # [C]
    layer_output = layer_fn(layer_input)    # [C]
    residual_mix = h_res @ x                # [n, C]
    write_back = h_post[:, None] * layer_output[None, :]
    return residual_mix + write_back

Original HC notation: depth and width connections

The original Hyper-Connections paper writes the connection matrix as:

For a network layer $\mathcal{T}$ and hyper-hidden matrix $\mathbf{H}\in\mathbb{R}^{n\times d}$:

Mapping this to the mHC notation:

HC further decomposes this into:

The intuition:

• depth-connections learn how much each stream should preserve old information versus accept new layer output;
• width-connections learn how streams communicate laterally at the same depth.

The original HC paper also introduces dynamic Hyper-Connections, where the maps are input-conditioned:

The small learned gates $s_\alpha,s_\beta$ keep the dynamic part small at initialization. In the original paper, dynamic HC with expansion $n=4$ was especially useful in language-model pretraining and achieved much faster convergence in the OLMoE setting.

Why HC can improve models

HC increases macro-architectural flexibility without making attention or MLP blocks wider.

That is the important scaling argument:

• increasing $C$ makes attention projections and MLPs much more expensive;
• increasing $n$ adds multiple residual streams, but the layer body still sees a $C$-dimensional vector;
• since typical $n$ is small, e.g. $n=4$, the extra coefficient maps are much cheaper than the block itself.

Conceptually, HC gives the model a learned routing topology across depth.

Special cases:

• if the maps are chosen one way, HC behaves like a standard sequential residual network;
• if chosen another way, it can emulate parallel transformer blocks;
• with dynamic maps, different tokens can use different soft layer arrangements.

This is why the HC paper argues that Hyper-Connections can learn mixtures of Pre-Norm-like, Post-Norm-like, sequential, and parallel wiring.

The instability: products of residual mixing matrices

The problem appears when we stack many HC layers. Expanding the single-layer update over depth gives:

This has two pieces:

1. The shallow feature $\mathbf{x}_l$ transformed by a product of residual mixing matrices.
2. A sum of every previous layer output, each also transformed by products of later residual mixing matrices.

With a standard residual connection, the shallow feature path is just $\mathbf{x}_l$. With HC, it is:

That product is the danger.

If each $\mathcal{H}_l^{\mathrm{res}}$ is unconstrained, even small deviations from identity compound. A simple scalar analogy shows the issue:

Tiny per-layer amplification becomes explosion; tiny per-layer attenuation becomes disappearance. Matrices are worse because different directions can amplify or shrink differently.

The mHC paper measures this using Amax Gain Magnitude:

For a composite residual map:

$G_{\mathrm{fwd}}$ captures worst-case forward signal gain, and $G_{\mathrm{bwd}}$ captures worst-case backward gradient gain.

DeepSeek reports that HC can hit Amax gain values near $3000$ in 27B-scale experiments. That is no longer a gentle residual highway; it is a learned multiplier chain.

mHC’s core fix: constrain residual mixing to the Birkhoff polytope

mHC keeps the useful part of HC: multiple residual streams and learned stream mixing.

But it constrains the residual mixing matrix:

where:

This is the set of doubly stochastic matrices, also called the Birkhoff polytope.

A matrix is doubly stochastic if:

1. all entries are non-negative;
2. every row sums to $1$;
3. every column sums to $1$.

Example:

is doubly stochastic. It mixes the two streams, but it cannot create or delete average mass.

Why doubly stochastic matrices stabilize depth

There are three important properties.

They preserve the all-ones direction

If $H\mathbf{1}=\mathbf{1}$, each row sums to $1$. If $\mathbf{1}^\top H=\mathbf{1}^\top$, each column sums to $1$.

So $H$ acts like a conservative mixing operator over streams. It can redistribute signal among streams, but it cannot globally scale the stream average up or down.

They are non-expansive in spectral norm

For a non-negative doubly stochastic matrix:

Using the norm inequality:

Therefore:

This directly attacks gradient explosion through $\mathcal{H}^{\mathrm{res}}$.

They are closed under multiplication

If $A$ and $B$ are doubly stochastic, then $AB$ is also doubly stochastic:

and $AB\ge 0$ because $A,B\ge 0$.

This is the crucial depth property. The composite map:

is also doubly stochastic if every factor is doubly stochastic. Stability survives depth.

Geometrically, the Birkhoff-von Neumann theorem says the Birkhoff polytope is the convex hull of permutation matrices:

So mHC residual mixing is a soft mixture of stream permutations. It can route and combine streams, but within a conservative envelope.

How mHC constructs the constrained maps

mHC first computes unconstrained pre-activations from the flattened $n$-stream state:

Then:

The final constrained maps are:

Two details matter:

• sigmoid makes $\mathcal{H}^{\mathrm{pre}}$ and $\mathcal{H}^{\mathrm{post}}$ non-negative, reducing cancellation from positive/negative coefficient compositions;
• the factor $2$ on $\mathcal{H}^{\mathrm{post}}$ centers the post map around $1$ when its pre-activation is near zero, because $2\sigma(0)=1$ - making the hyper connections behave exactly like standard residual connections at the beginning of training.

Sinkhorn-Knopp: projecting to a doubly stochastic matrix

Given an unconstrained matrix $\widetilde{H}$, mHC first makes it positive:

Then it alternates column and row normalization:

where:

Each column-normalization step fixes column sums. Each row-normalization step fixes row sums. Repeating the process converges toward a doubly stochastic matrix under standard positivity assumptions.

import numpy as np

def sinkhorn_knopp(logits, iters=20, eps=1e-12):
    """
    Project an unconstrained square matrix toward the Birkhoff polytope.

    logits: [n, n] unconstrained residual-mixing scores
    returns: approximately doubly stochastic [n, n] matrix
    """
    m = np.exp(logits)  # make entries positive

    for _ in range(iters):
        m = m / (m.sum(axis=0, keepdims=True) + eps)  # normalize columns
        m = m / (m.sum(axis=1, keepdims=True) + eps)  # normalize rows

    return m

In DeepSeek’s mHC experiments:

Because finite iterations are approximate, mHC does not produce a mathematically exact doubly stochastic matrix in practice. But the paper reports that the composite Amax gain stays bounded around $1.6$, compared with nearly $3000$ for vanilla HC.

The full mHC update

Once the maps are constrained, the layer update keeps the HC form:

But now:

The solution is not “go back to identity.” Identity would be stable but would remove residual-stream communication. mHC chooses the larger stable family of conservative mixing matrices.

The mHC layer is the same high-level update as HC, but the maps are constructed through constraints before the residual update is applied:

def mhc_maps(x, phi_pre, phi_post, phi_res, b_pre, b_post, b_res, alpha):
    """
    x: [n, C] residual streams
    Returns constrained h_pre, h_post, h_res.
    """
    flat = rms_norm(x.reshape(1, -1))

    pre_logits = alpha["pre"] * (flat @ phi_pre) + b_pre
    post_logits = alpha["post"] * (flat @ phi_post) + b_post
    res_logits = alpha["res"] * (flat @ phi_res).reshape(x.shape[0], x.shape[0]) + b_res

    h_pre = sigmoid(pre_logits)          # non-negative aggregation
    h_post = 2.0 * sigmoid(post_logits)  # starts near 1 when logits are near 0
    h_res = sinkhorn_knopp(res_logits)   # approximately doubly stochastic

    return h_pre.squeeze(), h_post.squeeze(), h_res

That is the core idea:

HC widened the residual highway into multiple lanes. mHC adds traffic rules so lane changes cannot create unbounded amplification.

Systems problem: HC is FLOP-light but I/O-heavy

HC and mHC keep attention/MLP FLOPs mostly unchanged, but they create a wider residual stream of size $nC$.

For a standard residual merge, the paper’s simplified per-token I/O is:

For HC, the residual-stream maintenance costs roughly:

For $n=4$, this is a serious memory-bandwidth problem even if FLOPs look cheap. The widened residual stream also increases activation storage and pipeline communication.

mHC’s efficiency optimizations

DeepSeek’s mHC paper adds infrastructure work to make the math practical.

Kernel fusion

mHC fuses coefficient generation, RMSNorm-related operations, Sinkhorn-Knopp, and residual merge paths where possible. In particular, it fuses the post/residual application:

The paper reports this reduces reads for that kernel from:

and writes from:

Recomputing instead of storing

Because the mHC coefficient kernels are cheaper than the attention/MLP block, the paper discards many intermediate mHC activations after the forward pass and recomputes them during backward.

For a recomputation block of $L_r$ consecutive layers, the memory objective is:

The first term is persistent storage for block starts. The second term is transient memory during recomputation.

The recomputation strategy is:

def backward_with_mhc_recompute(block_start_x, saved_layer_outputs, layers):
    """
    Store only block_start_x plus heavy layer outputs.
    Recreate cheap mHC maps and intermediate residual states during backward.
    """
    x = block_start_x
    states = [x]

    for layer in layers:
        h_pre, h_post, h_res = mhc_maps(x, *layer.mhc_params)
        x = h_res @ x + h_post[:, None] * saved_layer_outputs[layer.id][None, :]
        states.append(x)

    # Backprop now uses recomputed states instead of storing all of them.
    return run_backward_from_recomputed_states(states, layers)

Pipeline overlap

For large-scale distributed training, the $n$-stream state increases pipeline communication. DeepSeek extends its DualPipe schedule so mHC communication and recomputation can overlap with useful compute, especially around pipeline-stage boundaries.

The result reported for $n=4$ is only 6.7% additional training time overhead, after these optimizations.

Empirical picture

The original HC paper reports large gains in OLMo/OLMoE-style pretraining. In the OLMoE-1B-7B setting, DHC with $n=4$ converged about 1.8x faster than the baseline and improved several downstream metrics.

The mHC paper’s 27B results preserve most of the HC benefit while stabilizing training. The reported benchmark table is:

The stability story is more important than the raw score table:

• vanilla HC can outperform early, but the gradient norm and loss can surge at scale;
• mHC keeps the multi-stream expressivity while making the residual product behave like conservative mixing;
• the composite gain drops from nearly $3000$ in HC to a bounded value around $1.6$ in mHC’s finite-iteration implementation.

Why the constraint works

The key mechanism in mHC is the doubly stochastic constraint on $\mathcal{H}^{\mathrm{res}}$. Its effect is easiest to see through the residual product across depth:

1. The instability is not mainly about the immediate single-layer output. It is about the product of residual mixing matrices across many layers.
2. Row sums and column sums have different interpretations: row sums correspond to forward signal gain; column sums correspond to backward gradient gain.
3. Doubly stochastic matrices are stable because of three linked facts: non-negativity, row/column conservation, and closure under multiplication.
4. The constraint needs systems support. Without kernel fusion, recomputation, and pipeline overlap, the widened residual stream would be too I/O-heavy.
5. Sinkhorn-Knopp is approximate at finite iteration count. Later variants such as mHC-lite and KromHC target exact or cheaper parameterizations of doubly stochastic residual maps.

Open engineering questions

mHC stabilizes the residual product, but it leaves several practical trade-offs.

Finite Sinkhorn iterations

With finite $t_{\max}$, Sinkhorn-Knopp only approximately reaches the Birkhoff polytope. The mHC paper uses $t_{\max}=20$, which is effective empirically, but still leaves a small approximation gap.

The mHC-lite paper proposes constructing doubly stochastic matrices directly as convex combinations of permutation matrices. This uses the Birkhoff-von Neumann theorem and avoids iterative Sinkhorn projection, but it introduces its own parameterization and engineering trade-offs.

Parameter cost of residual maps

In mHC, the residual projection uses:

so the residual-map parameterization scales like:

This is manageable for small $n$, but KromHC explores Kronecker-product residual matrices to reduce complexity while preserving exact double stochasticity.

Expressivity versus constraints

Doubly stochastic matrices forbid negative mixing. This supports conservative signal propagation, but may restrict expressivity. Later work explores alternative manifolds with different stability constraints. The design pressure is the same: make the residual stream learnable without allowing depthwise composition to create pathological gain.

Summary

Hyper-Connections generalize residual connections by turning the single residual stream into $n$ streams and learning how to aggregate, update, and mix them:

This gives the model a richer topology across depth. But the same flexibility creates instability because:

can amplify or attenuate signals exponentially.

mHC solves this by forcing:

The result is a residual mixer that can exchange information across streams, but whose products remain conservative and stable. The real contribution is the combination of topology, math, and systems work: widen the residual stream, constrain the dangerous map, then make the implementation bandwidth-aware.

Sources

• Defa Zhu et al., Hyper-Connections, arXiv:2409.19606.
• Zhenda Xie et al., mHC: Manifold-Constrained Hyper-Connections, arXiv:2512.24880.
• How mHC Reinvents Residual Connections, source video for the figures used in this explainer.
• Yongyi Yang and Jianyang Gao, mHC-lite: You Don’t Need 20 Sinkhorn-Knopp Iterations, arXiv:2601.05732.
• Wuyang Zhou et al., KromHC: Manifold-Constrained Hyper-Connections with Kronecker-Product Residual Matrices, arXiv:2601.21579.

Follow me on Twitter, Github or connect on LinkedIn.

Introduction

A scan (prefix sum) is a deceptively small primitive: given an input array X, produce an output array Y where Y[i] = X[0] + X[1] + … + X[i]. On the GPU this is hard to do efficiently because each output depends on all previous elements, which sounds serial. The kernels in this repository explore multiple ways to restructure this computation so that thousands of threads can participate without breaking correctness.

There are two broad families here:

1. Hierarchical (multi-pass) scans: scan within blocks, scan the block totals, then redistribute those totals back into the output. This is the most standard GPU scan strategy and maps cleanly to CUDA’s execution model.
2. Single-pass scans: attempt to compute the full array scan in a single kernel launch using inter-block coordination (domino propagation or decoupled lookbacks). These are more complex but avoid extra kernel launches.

A CUB baseline in src/cub_scan.cu uses cub::DeviceScan::InclusiveSum as a reference for performance and correctness.

Quick CUDA primer (for context)

If you are new to CUDA, three concepts show up repeatedly in the kernels below:

• Warps: threads are executed in groups of 32. Many performance optimizations (like warp shuffles) are designed around this unit.
• Shared memory: fast, on-chip memory shared by threads in a block. Most scan algorithms use shared memory for their per-block scan stages.
• Synchronization: __syncthreads() synchronizes threads within a block. There is no built-in global synchronization across blocks inside a kernel, which is why single-pass scans must use explicit memory protocols to coordinate. __threadfence() orders a thread’s global memory writes so they are visible to other threads/blocks on the device before it continues.

You’ll also see the term coalesced memory access. A global memory access is coalesced when consecutive threads in a warp access consecutive addresses, allowing the GPU to serve the warp with fewer memory transactions. This is a major performance factor, and it strongly influences how the kernels index into memory.

Hierarchical Scan Algorithms

Idea of hierarchical scan

Hierarchical scan decomposes the full scan into three stages that are easy to parallelize:

1. Per-block scan: each block scans a contiguous chunk of the input and writes the prefix results for that chunk into the output array. Each block also emits one number: the block total (sum of its entire chunk).
2. Scan the block totals: scan the array of block totals to produce a prefix over blocks. Block 0 adds 0; block \(b>0\) adds the scanned total of block \(b-1\) (i.e., the sum of everything before block \(b\)). If this array is long, it is scanned hierarchically using multiple levels.
3. Redistribution (add carry‑in): each block adds its carry‑in to every element of its local output, turning a block-local prefix into a correct global prefix.

In the source, the “block totals” buffer is called partialSums: it is an auxiliary array with one entry per thread block.

Concretely, after stage 1 each block has computed the right prefix order relative to the start of its own chunk, but every block except block 0 is missing a constant offset (the sum of all earlier chunks). Scanning the block totals computes exactly those offsets.

Example with BLOCK_SIZE = 4:

Input:        [1 2 3 4 | 5 6 7 8]
Local scans:  [1 3 6 10 | 5 11 18 26]
Block totals: [10, 26]
Scan totals:  [10, 36]
Add carry-in: [1 3 6 10 | (5+10) (11+10) (18+10) (26+10)]
            = [1 3 6 10 | 15 21 28 36]

This structure is consistent across:

• src/hierarchical_kogge_stone*.cu
• src/hierarchical_brent_kung*.cu
• src/hierarchical_warp_tiled*_optimized.cu

A key idea for readers: the per-block scan only handles a local segment. The global correctness comes from the second and third stages, which propagate block totals across the array.

When the block-totals scan needs multiple levels

Stage 2 scans one value per block. If each block handles \(B =\) BLOCK_SIZE input elements, then the number of blocks is \(M = \lceil N / B \rceil\), so the block-totals array has length \(M\).

A single CUDA block in these kernels scans at most \(B\) values (one value per thread in shared memory), so stage 2 is:

• one-block when \(M \le B\) (equivalently \(N \le B^2\); for \(B = 1024\), about one million elements),
• a small recursive hierarchy when \(M > B\).

Conceptually, you build a short “pyramid” of group totals:

• Level 0: per-block totals (length \(M_0 = M\))
• Level 1: totals of contiguous groups of \(B\) entries from level 0 (length \(M_1 = \lceil M_0 / B \rceil\))
• …
• stop at the first level \(L\) with \(M_L \le B\)

Then you run the same up/down structure across levels:

1. Up-sweep: scan each level in block-sized segments and write each segment’s total into the next level.
2. Top scan: scan the final level in one block.
3. Down-sweep: propagate prefixes back down by adding the scanned prefix of earlier segments (the carry‑in) into every element of the lower level.

In src/hierarchical_kogge_stone.cu, this logic is packaged as ScanLevels: level 0 is the block-totals buffer (partialSums in the code), and higher levels are temporary allocations that shrink by ~1024× each step:

while (curr_len > BLOCK_SIZE) {
    unsigned int next_len = cdiv(curr_len, BLOCK_SIZE);
    T* sums_d = nullptr;
    CHECK_CUDA_ERROR(cudaMalloc(&sums_d, next_len * sizeof(T)));
    levels.data.push_back(sums_d);
    levels.lengths.push_back(next_len);
    curr_len = next_len;
}

The scan then follows the “up-sweep / top / down-sweep” pattern literally:

for (size_t level = 0; level + 1 < levels.data.size(); ++level) {
    unsigned int len = levels.lengths[level];
    dim3 gridSize(cdiv(len, BLOCK_SIZE));
    kogge_stone_segmented_scan_kernel<<<gridSize, blockSize>>>(
        levels.data[level], levels.data[level], levels.data[level + 1], len);
}

kogge_stone_scan_kernel<<<dim3(1), blockSize>>>(
    levels.data.back(), levels.lengths.back());

for (int level = static_cast<int>(levels.data.size()) - 2; level >= 0; --level) {
    unsigned int len = levels.lengths[level];
    dim3 gridSize(cdiv(len, BLOCK_SIZE));
    redistribute_sum<<<gridSize, blockSize>>>(
        levels.data[level], levels.data[level + 1], len);
}