GPU Memory System Optimization for LLM Inference

MoH-guided head-wise offloading: routing traces, head regularity, and the transfer cost of acting on them

01Motivation

LLM inference is bounded by attention-related memory cost: the KV cache footprint, the traffic it generates, and the host-device transfers that follow when it does not fit. Mixture-of-Head attention already produces a per-token signal about which heads matter. If that signal is stable enough, it can drive a memory decision — keep these heads resident, move those to host memory — instead of only a compute decision. This project tests whether that holds, and what it costs when it does.

Key question Can MoH routing signals reduce effective attention cost, by using fewer heads or treating heads differently, without hurting output quality?

02System problem

The target is a Mixture-of-Head attention model served with PyTorch and Transformers, with hooks on the router output that log which heads each token selects. Attention state lives in GPU memory; host DRAM is far larger but sits behind PCIe. Any head-wise offloading plan therefore has two costs that do not appear in a head count: the transfer of head state across PCIe, and the stalls when that transfer is not overlapped with compute. The router trace tells us which heads a policy would keep; the transfer measurement tells us whether keeping them is worth anything.

Where the cost sits. Attention state is resident in GPU memory, host DRAM is large but a PCIe crossing away, and every head-wise offloading decision pays for that crossing whether or not it saves compute.

03Approach

Token by head selection grid recorded from the router. If neighbouring tokens select the same heads, a residency policy can reuse them; the Jaccard similarity between adjacent columns is how that regularity is measured rather than assumed.

1Collect router tracesRun inference on controlled prompts with scripts/run_router_trace.py and log per-token head selection and router scores to a trace file.

2Characterise the tracesPlot head activation frequency and Jaccard similarity between neighbouring tokens and across layers, to test whether selection is regular enough for a residency policy to predict.

3Sweep rhoVary the fraction of heads kept and measure output quality against it, then extend the same sweep to latency, throughput, peak GPU memory, PCIe bytes, and overlap ratio.

4Compare policiesEvaluate naive offloading, keep-hot with a window tau, and prefetch with overlap, recording for each what it has to know in advance and where it loses.

04Current findings

Accuracy vs rho

There is a gentle regime where dropping heads costs little quality, and a threshold past which the loss becomes sharp. The useful question is where that knee sits, not whether it exists.

Regularity

Head selection is less regular than the offloading idea needs. Similarity between neighbouring tokens is not high enough to justify a simple reuse heuristic, so a policy built on "the same heads stay hot" would misfire more often than the trace plots first suggested.

Transfers

Fewer active heads does not mean less time. Whether a head-wise policy wins depends on transfer size, pinned memory, and overlap, so the system metrics have to be measured alongside quality rather than after it.

The same two layers on two timelines. Serial: the transfer for layer i+1 starts after layer i finishes computing. Prefetch: the transfer runs underneath the compute, and only the uncovered part still costs time.

Portfolio note Portfolio note: this page is not a paper. It shows the question, the method, and the judgement calls as they stand today, and it is rewritten as results arrive.

05What went wrong

Symptom

Reducing the number of heads, or introducing offloading, sometimes produced almost no gain and sometimes made the run slower than the baseline.

Hypotheses

Many small transfers dominate the latency; or pageable memory and synchronous copies prevent any overlap with compute; or the router tracing itself is heavy enough to contaminate the timings.

Tests

Microbenchmarked host-to-device and device-to-host bandwidth against transfer size, and copy-compute overlap separately; then ran ablations that toggled logging, batched the transfers, and separated trace collection from the timed run.

Fix

Pinned memory with asynchronous transfers, a coarser transfer granularity, keep-hot with a window tau, and trace collection moved out of any run whose time is being reported.

Takeaway

Algorithmic sparsity is not a system result. Transfer granularity, pinning, and overlap decide whether a smaller head count turns into a smaller number of milliseconds.

06Next steps

Add system metrics to every sweep — latency, throughput, peak GPU memory, PCIe bytes moved, overlap ratio — so a quality loss can be traded against a measured gain instead of an assumed one.
Measure host-to-device and device-to-host bandwidth against transfer size on the target GPU, and set the offloading granularity from that curve rather than from the head layout.
Run keep-hot and prefetch against a no-offload baseline on identical prompts, and report the prompts where the policies lose as well as the ones where they win.