llama : keep explicit host buffer overrides when using mmap - #28223
llama : keep explicit host buffer overrides when using mmap#28223Inovello wants to merge 2 commits into
Conversation
37534fa to
6b49f41
Compare
Testing revealed that a host buffer type, specifically CUDA_Host was being replaced with CPU buffer type under mmap. Patch makes it so that when the buffer type comes from an explicit -ot rule, it skips the safety check it does (The loader's rule that essentially replaces a host buffer type with a plain CPU one when a model is memory mapped) and keeps the type you asked for which enables CPU resident experts to live in the pinned memory for op offload while the rest of the model stays memory mapped. Prefill of a 26k prompt went from 166 t/s to 379 t/s on 2 x RTX 3090 with 40 expert layers on the host. The change lives in the check in create_tensor() in src/llama-model-loader.cpp, which now skips the host buffer downgrade when the new buft_overridden flag is set, and in parse_tensor_buffer_overrides() in common/arg.cpp, which now accepts CUDA_Host as a target. Assisted-by: Claude Fable 5.1
6b49f41 to
5cfa6a8
Compare
1. Issue and cause: Load time was incredibly long and it was caused by load_all_data copying each expert tensor from the mmap into the pinned buffer. The mmap carries POSIX_MADV_RANDOM under --numa distribute, so the memcpy faulted one 4 KiB page at a time with no readahead. 2. What changed: When the destination buffer is the host memory, the tensor will be read from the file with read_raw instead of being copied from the mmap. llama-mmap.cpp and the normal mmap path are untouched. 3. Result: Cold load went from 512 s to 168 s on 2 x RTX 3090 with 101.75 GiB of experts on CUDA_Host; the copy phase went from 450 s to 104 s. This doesn't affect the prefill or the output, just the load times. Assisted-by: Claude Fable 5.1
|
Whilst experimenting with ways to circumvent the load time cost, I discovered that the reason for the pinned load being slow is not the pinning. Out of the 490-512 s, the pinned allocation only takes ~60 s whilst the copy takes ~450 s. It turns out the copy reads the tensors through mmap one 4 KiB page fault at a time because --numa distribute puts POSIX_MADV_RANDOM on the mapping and readahead is off. The fix for this is 9 lines in load_all_data(), in the same file this PR originally touches. When the destination buffer is host memory (i.e. CUDA_Host), it will now read the tensor from the file with read_raw instead of copying through the mmap one. llama-mmap.cpp remains as is, so the normal mmap NUMA behaviour is untouched. Measured on 2x RTX 3090 and 2x E5-2696 v4 with Qwen3.8-Flash-Next UD-Q6_K_XL, all 48 expert layers (101.75 GiB) overridden to CUDA_Host, --numa distribute. Both cold loads within the same hour:
The change was written with an agent's help and verified by me on the box above, same as the first commit. |
|
Independent A/B on a different box (from the r/LocalLLaMA thread, posted with the tester's permission) I was lucky to have a Reddit user (u/Beneficial-Ad-8127) run the first commit of this PR (5cfa6a8, host overrides kept under mmap) on hardware I do not have. Their numbers, unchanged, medians of 3: Setup: RTX 5090 32 GB, Ryzen 7 9800X3D, 192 GB DDR5, Windows. Qwen3.8-Flash-Next 4.27 bpw Q4_K_M (33-shard GGUF), mmap on,
So 2.3-2.8x prefill at every depth, decode flat, and no load-time penalty on Windows (the ~48 GB page-lock went through without error). Three notes:
|
1 participant
Overview
Testing revealed two things here that stop -ot "...=CUDA_Host" from working with a model that's mmapped. First thing is that on master, the argument parser does not offer host buffer types as -ot targets at all, so the flag is rejected at startup ("Available buffer types: CPU, CUDA0") and second, even when a host buffer type reaches the loader, it is replaced with a normal CPU buffer type under mmap (the loader's safety rule that avoids a second copy of bytes that are already in the mapping). The patch registers the device's host buffer types as -ot targets and also makes it so that when the buffer type comes from an explicit -ot rule, it skips the safety check it does and keeps the type you asked for which enables CPU resident experts to live in the pinned memory for op offload while the rest of the model stays memory mapped. Prefill of a 26k prompt went from 166 t/s to 330 t/s with a cold page cache, and to 379 t/s once the per-layer embedding rows were cached (warm), on 2x RTX 3090 with 40 expert layers on the host.
Now as for why the safety check exists and behaves the way it does. Under mmap, the bytes are already stored in the memory through the mapping, so putting a tensor in a host buffer makes a second copy. For the automatic choice that is waste, so the loader downgrades it, and that behaviour is unchanged. The change this PR brings only applies when the user names a host buffer type in an -ot rule. Default loading, --cpu-moe, --n-cpu-moe are not affected so that means the extra copy is the cost the user willingly chooses to pay. If the pinned allocation fails, the buffer falls back to pageable memory. A visible warning can be added if needed.
To test it out, you need to use the -ot rule with the target as a CUDA_Host, for example
-ot "ffn_(gate|up|down)_exps\.weight=CUDA_Host", and on a Dual CPU machine (Like Dual Xeons), you run under numactl --interleave=all.The change lives in the check in create_tensor() in src/llama-model-loader.cpp, which now skips the host buffer downgrade when the new buft_overridden flag is set, and in parse_tensor_buffer_overrides() in common/arg.cpp, which now accepts the device's host buffer types (e.g. CUDA_Host) as targets.
The second commit (c59754b) fixes the load time of that pinned buffer. Under --numa distribute the mapping is MADV_RANDOM, so copying the expert tensors from it into the pinned buffer went through page faults one 4 KiB page at a time (236 MB/s, 450 s of a 512 s load). Host destination tensors are now read from the file with read_raw instead. Load to ready went from 512 s to 168 s on the same box, prefill and greedy output unchanged, the buffer still remains pinned, and the plain mmap path (no host override) is untouched. Details and the full table are in my comment below.
Additional information
Setup: Qwen3.8-Flash-Next UD-Q6_K_XL (177B MoE, 512 experts, 10 active), 2x RTX 3090 24 GB (PCIe 3.0 x16),
2x Xeon E5-2696 v4, 188 GiB DDR4-2133, Ubuntu 24.04, CUDA 12.0. Context 261,888,
-ub 2048 -b 4096, f16 KV,4 expert layers per GPU, the other 40 layers (85 GiB) in CPU RAM. Server timings from
/completion, temperature 0.7,page cache dropped before each load,
--cache-ram 0.CPU(mmap, pageable)CUDA_Host(this PR)"PLE cold/warm": The model reads the per layer embedding table lazily. The first prompt reads the rows from disk, a second prompt with the same text reads the rows in the page cache. The warm-cache pass for the mmap row was not measured because my box lost power during that run.
Now as the table shows, loading time is longer because mmap defers reading until pages are touched, while a host buffer is allocated and filled at load time, so all 85 GiB are read from disk and page-locked before the server is ready. With the first commit alone that was 96 s to 488 s, most of it the page fault copy described above; with the second commit it is 96 s to 168 s. It is paid only by users who opt in, in exchange for 2x to 2.3x prefill on a server that runs for days.
(a) The run was cut off at 446 of 512 tokens by a reboot. mmap decode depends on NUMA page placement.
Also worth stating: ik_llama.cpp (fused MoE + pinned) does 353 t/s prefill on the same box, but 10 t/s decode.
Here are some related PR's/Issues I found: #26659, #26110, #25859, #26448, #28136.
Requirements