Fixing ollama's MoE memory estimate — and a 3× kernel speedup that wasn't

Ollama and llama.cpp predict a model’s GPU memory footprint during startup to determine how many transformer layers can be placed into GPU VRAMThe GPU's onboard memory, where the weights, KV cache, and compute buffer all have to fit.. If the estimator underestimates memory consumption, the runtime either forces unnecessary layer offloading to CPU or triggers out-of-memory (OOM) allocation failures during runtime.

This post analyzes Ollama’s memory estimation accounting alongside GPU kernel execution in llama.cpp. A rewrite of Ollama’s MoE memory estimator (PR #17201) corrected a -20.6% memory under-estimation down to -6.0%. Additionally, we examine an initial 3.43× speedup benchmark in llama.cpp’s MMQ kernel that failed numerical correctness checks on large matrix shapes.

Auditing Ollama’s MoE Memory Footprint Estimator

When placing model layers between GPU VRAM and system RAM, Ollama calculates expected allocations per layer. For mixture-of-expertsA model whose layers hold many expert sub-networks and fire only a few per token, so its footprint is dominated by weights that differ from layer to layer. (MoE) architectures, legacy allocation estimates ran ~20% below actual memory allocations.

Comparing estimated allocations against actual llama-server runtime allocations revealed two primary sources of error:

1. KV Cache Sizing: The KV cacheThe key and value vectors stored for every token, sized by layers, KV-head count, head dimension, and context length. was calculated using a single averaged head dimension across all layers.
2. Compute Scratch Buffer: Memory allocated for intermediate tensor execution (compute bufferScratch space llama.cpp reserves for the largest intermediate tensors in one graph step; it grows with context and batch.) was modeled as zero.

Using a single averaged head dimension fails on modern architectures. DeepSeek-V2’s Multi-Head Latent Attention (MLA) specifies asymmetric key and value dimensions (key_length = 192, value_length = 128); averaging these dimensions under-sized the KV cache by ~20%. Similarly, Grouped-Query Attention (GQAGrouped-query attention: several query heads share one key/value head, so the KV cache is far smaller than the query-head count suggests.) shares key/value heads across query groups; averaging head counts obscured this ratio.

PR #17201 replaced global averaging with per-layer metadata extraction, reading exact key/value dimensions and KV-head counts per layer while assigning a non-zero lower bound to compute scratch buffers.

Evaluating llama.cpp MMQ Kernel Tile Sweeps and Correctness Validation

llama.cpp’s quantized matrix multiplication kernel (MMQ) uses a fixed default tile configThe compile-time tuple that sets how a kernel splits a matmul into tiles across threads, registers, and VRAM. (nthreads=256, occupancy=1, I=128, K_vram=256, stream_k=true) across multiple NVIDIA GPU generations (A100, RTX 3090, RTX 4090, H100).

We evaluated candidate tile configurations to test whether hardware-specific tuning could improve throughput. Increasing VRAM staging tile size to K_vram=512 produced a 3.43× throughput jump in microbenchmarks (138.9 TFLOP/s vs 40.5 TFLOP/s on Q4_K) and a 1.60× speedup in llama-bench.

However, validating the candidate configuration against llama.cpp’s test-backend-ops suite revealed numerical failures. The K_vram=512 tile configuration failed 25 out of 78 correctness tests for Q4_K and Q6_K tensor shapes, returning Normalized Mean-Squared Error (NMSENormalized mean-squared error, MSE divided by the variance of the true output. Zero is exact; 1.0 means the output is uncorrelated with the correct answer.) values near 1.0 on large-matrix prefill shapes.

Benchmarking INT8 k-Quant Kernels vs. FP16 cuBLAS

llama.cpp provides INT8 quantization paths for Q4_K and Q6_K weights (k-quant). Upstream documentation notes that at larger batch sizes (e.g. batch $\ge 512$), FP16 cuBLAS GEMMGeneral matrix multiply, the dense linear-algebra operation at the core of every transformer layer. execution can outperform custom INT8 dequantization kernels once workload intensity becomes compute-bound.

We benchmarked INT8 k-quant against FP16 cuBLAS on a Qwen2.5-1.5B model using an RTX 3090 GPU:

On smaller models or memory-bound batch sizes (batch size 1 decode), the INT8 k-quant path maintains a clear advantage (1.74× speedup over FP16 cuBLAS) by reducing HBM weight transfer volume.

DeepSeek Multi-Head Latent Attention (MLA) Execution Paths

DeepSeek architectures utilize Multi-Head Latent Attention (MLA) to compress key and value projections into latent representations. We analyzed MLA execution paths within llama.cpp’s FlashAttention implementation (fattn.cu).

MLA operates in two structural representations within llama.cpp:

Next Steps

1. Gated MMQ Tile Sweeps: Re-evaluate MMQ tile parameters across GPU microarchitectures with inline test-backend-ops correctness validation.
2. Compute-Bound Crossover Benchmarking: Benchmark INT8 k-quant kernels against FP16 cuBLAS on 30B+ parameter models at batch sizes $\ge 512$.
3. MLA Fallback Optimization: Optimize memory double-buffering for uncompressed MLA fallback execution paths in fattn.cu.

Conclusion

Systematic profiling of memory accounting and kernel paths ensures predictable deployment behavior. Correcting Ollama’s MoE memory estimator (PR #17201) prevents premature CPU offloading and OOM failures, while combining throughput benchmarks with numerical correctness gates prevents false performance conclusions during kernel optimization.

Glossary