PR #16222: ml/ggml: add OLLAMA_PINNED_HOST_BUFFER opt-in for faster prefill on partial-offload models

• Repository: ollama/ollama
• Author: lingyuncai
• State: open | merged: False
• Link: https://github.com/ollama/ollama/pull/16222

Description (problem / solution / changelog)

Summary

Adds OLLAMA_PINNED_HOST_BUFFER=1 opt-in to route CPU-resident model weights through CUDA pinned host memory (cudaMallocHost) instead of plain pageable memory, closing the prefill gap described in issue #16221 .

Motivation

On partial-offload setups (model size > VRAM), ollama runner's CPU buffer type list is built from ggml_backend_dev_buffer_type() only — plain malloc-backed pageable memory. GPU devices are skipped entirely in that loop, so the pinned host buffer type exposed via ggml_backend_dev_host_buffer_type() is never used for weight allocation.

llama runner's make_cpu_buft_list explicitly prepends the first GPU's host buffer type to the list for exactly this reason (see comment in llama-model.cpp:336-350). This means CPU-resident weights in llama runner land in cudaMallocHost-backed pinned memory, while ollama runner's go through pageable memory with a driver staging copy on every H2D transfer. With ~28 GiB of CPU-resident weights per prefill batch, this is a meaningful per-batch cost.

Implementation

Backend.New() builds the CPU buft list in three explicit passes, matching llama.cpp's make_cpu_buft_list order (ACCEL → GPU host → CPU):

// ACCEL pass (unchanged)
for _, d := range accels {
    bt := C.ggml_backend_dev_buffer_type(d)
    cpuDeviceBufferType.bts = append(cpuDeviceBufferType.bts, bt)
    btDeviceMemory[bt] = &requiredMemory.CPU
}

// GPU host pass (opt-in)
if envconfig.PinnedHostBuffer() {
    for _, d := range gpus {
        hostBuft := C.ggml_backend_dev_host_buffer_type(d)
        if hostBuft == nil {
            continue
        }
        cpuDeviceBufferType.bts = append(cpuDeviceBufferType.bts, hostBuft)
        btDeviceMemory[hostBuft] = &requiredMemory.CPU
        break
    }
}

// CPU pass (unchanged)
for _, d := range cpus {
    bt := C.ggml_backend_dev_buffer_type(d)
    cpuDeviceBufferType.bts = append(cpuDeviceBufferType.bts, bt)
    btDeviceMemory[bt] = &requiredMemory.CPU
}

createTensor already walks bts in order, so CPU-resident weights automatically land in pinned memory when the opt-in is set. No changes to op scheduling, decode path, or KV cache placement.

Results

Test platform:

• -29% prefill latency / +40% prefill tps vs default ollama runner

Note

Allocation failure fallback: if cudaMallocHost fails, createTensor advances to the next buft in the list — plain CPU buffer. Worst case is identical to the disabled path.

No-op cases:

• OLLAMA_PINNED_HOST_BUFFER unset (default)
• No CUDA GPU present (gpus slice is empty)
• Models that fully fit in VRAM (no CPU-resident weights to transfer)

Changed files

• envconfig/config.go (modified, +6/-0)
• ml/backend/ggml/ggml.go (modified, +38/-9)