qwen4exp: direct reads for the lazy PLE table (>2x prefill performance improvement on GB10) - #28136
qwen4exp: direct reads for the lazy PLE table (>2x prefill performance improvement on GB10)#28136coder543 wants to merge 2 commits into
Conversation
|
I independently tested The measurements below use a build with instrumentation-only changes to record With the ordinary lazy mmap path, the cold PLE work reproduced the behavior I had previously localized independently: 207,913 major faults occurred inside the PLE With the direct-read path, PLE-related major faults were effectively eliminated: 43 process-wide major faults with one direct-read worker and 77 with the PR-default worker count, versus ~208k with mmap. I also ran the direct reader with a single worker to separate the explicit-read path itself from parallelism:
The one-worker direct path was only ~1.15x faster than the serialized mmap path, while the PR-default 32 dedicated read workers reduced the stage wall by a further ~6.6x. Per-row read service time rose from 52.5 µs at one worker to 211.2 µs at 32, but aggregate read service divided by stage wall corresponds to ~26x effective read concurrency. So on this system the dominant gain comes from exposing substantial I/O concurrency, despite higher per-read latency with many concurrent read workers. Note the counters are not 1:1 — the mmap figure counts major faults while the direct figures count row reads, and the PR dedups per ubatch rather than globally. As a separate causal check with a different implementation, preloading the exact PLE hot set (208,771 pages / 815.5 MiB, derived from the gathered row indices) eliminated 207,488 / 207,488 PLE-loop major faults and reduced the same PLE-local cold cost from 17.9 s to 0.35 s, against a warm floor of ~0.16 s. That independently supports cold sparse PLE backing acquisition as the bottleneck. I did not observe a measurable PLE-local warm regression with the PR-default configuration (~154.6 ms direct vs ~154.0 ms mmap in this test; the measurement boundaries are not exactly identical). End-to-end prompt time is reported only as a reference on this machine because the 72.4 GB model substantially exceeds 16 GB VRAM and whole-request timing is highly sensitive to unrelated model/page-cache residency. |
|
Independent test on an unusual but real low-end serving config: 2× Xeon E5-2620 (Sandy Bridge, no AVX2/BMI2/FMA), 15 GB RAM, and 8 GPUs (2× RTX 3090 + 6× CMP 90HX) all behind PCIe Gen2 x4. The model is Qwen3.8-Flash-Next I patched this PR onto my tree and compared
That is ~+37% cold and ~+25% warm prefill on this box — smaller than the >2x seen on GB10, which I'd attribute to the fast NVMe already masking part of the demand-paging cost, but very consistent across repeats. Decode at 60K context is unchanged within noise (15.7–19.2 t/s in both modes). Two notes from porting:
Adopting |
|
mmap is a good default for streaming a file in once, and a poor fit for a scattered gather of tiny rows: a full page faulted to serve about ninety bytes, readahead waste on top, and a synchronous fault that drives the device at QD1 whatever it is capable of. Explicit reads from several workers fix both the amplification and the concurrency. Unified memory is the worst case here, since the weights leave nothing for the page cache and the rows are effectively always cold. It is going to matter on Apple silicon too, for the same reason. On a machine with spare RAM the alternative is to keep the table resident on the host, which gives full memory bandwidth and immunity to another model evicting the cache, but that option disappears as soon as the table cannot be resident. |
|
Confirmed from the "table cannot be resident" side of that fork: this box has 15 GB of host RAM with an 84 GB model mapped, so the page cache holds ~11 GB total and gets churned the moment anything else touches the disk — the resident-table option is off the table by construction. The +37% cold / +25% warm prefill I posted above is exactly that regime, and the numbers were stable across repeats. The QD1 point also matches what I saw: cold prefill on |
|
Strix Halo (gfx1151) llama.cpp build 10743 Model Method. Only Per rep: drop_caches → start server → send prompt (cold) → send again (warm) → kill. Cold prefill, tok/s (mean ± stdev, n=5)
Warm is a no-op -0.4% to +2.1% across all 12 cells, as expected once the rows are in page What changed (per cold prefill, mean, UD-Q5_K_XL; Q4 identical pattern)
|
|
If I understood correctly, this PR's new Since these tables are accessed via direct reads into temporary, short-lived working buffers rather than being mapped into virtual memory, they never accumulate in the OS page cache. This keeps the persistent memory footprint limited strictly to the model's main weights and KV cache. If so, loading a model with This is gold! |
|
If we end up needing this, I would like to see some cleaner way to implement this. Not sure what exactly, but the proposed change would not scale well with more models that might need this in the future. Even Gemma today should benefit from this. |
|
I had to add the following line to #include <string>am I the only one who had this issue? I am on windows Also Is windows support coming? |
ggml-org#28136 was written against a master where the PLE tensor was required, giving an outer-scope `const auto & ple_w = ml.require_weight(...)` that its direct-read path uses for the file index and offset. Master has since made the tensor optional (`if (const auto * ple_w = ml.get_weight(...))`, block-scoped pointer), so the merge is textually clean but does not compile: 'ple_w' was not declared in this scope. Hoist the lookup to outer scope as a pointer and guard the direct-read path on it. A metadata-only model has no file to pread from, so skipping direct reads there is the correct behaviour, not just the compiling one. Semantic conflict, not a textual one, so rerere cannot replay it -- this commit has to be re-applied by hand each time release is rebuilt.
|
@coder543 I would like to reproduce the tests on GB10. Could you please specify the exact model used (which quant?) and the full command line? |
|
I did my best to make it compile under Windows, but I don't have a Windows machine available at the moment. I can push a fix for that later. I didn't realize that MSVC would have a problem with If @ggerganov or someone else wants to suggest a better implementation, I would be happy to take a look at that, but otherwise, this PR seems like a real performance win and its hopefully fairly unobtrusive since it is very limited in scope. For Gemma 4 E2B/E4B, the access pattern seems to benefit a lot less from direct reads since it is returning ~5KiB per access, instead of randomly accessing very small values the way that For consistency, it would be a nice option to have on Gemma, but the benefit mostly seems to be exclusive to @eiffel31 the model that I was testing is this one. The command line is nothing special, what matters is setting |
Each cold PLE row demand-faults a 4 KiB page for ~90 bytes of data, capping cold diverse-text prefill at 218-360 tok/s vs ~785 warm on GB10. All n-gram row indices of a ubatch are known host-side before the graph runs, so under the new LLAMA_LAZY_MODE_DIRECT (--lazy-mode on-direct) they are staged into an input tensor with sorted, deduplicated, parallel pread()s and dequantized exactly like ggml_get_rows; downstream kernels unchanged, table stays on disk. Cold diverse prefill on qwen3.8-flash-next: 542-741 tok/s (2.0-3.1x, within ~6% of warm); warm, decode and greedy outputs bit-identical.
Move the row reader into arch-agnostic src/llama-lazy-reader.h, owned by llama_model_base keyed by tensor name; arches opt in with one load_lazy_reader call. qwen4exp now uses the shared reader, and gemma4 stages its per-layer rows the same way. Gemma's ~5 KiB contiguous rows show no prefill win either way, which is why the reader stays opt-in per tensor rather than unconditional. Also fix the missing <string> include in llama-mmap.h (MSVC needs it for the name() declaration) and make the help text generic so it does not stale when arches opt in.
|
I have pushed a commit showing what this looks like if refactored to support gemma4 as well. |
FYI Windows ISOs are downloadable for free and can be installed in a VM in <20 minutes and <30GB, the only "limitation" is a watermark on desktop, which you can ignore indefinitely. |
|
@coder543 I implemented the Windows side of the direct-read path and measured it — the port works and is correct, but on this box it is 2–4% slower than the mmap path, not faster. Patch and numbers below; the honest summary is that the win looks platform-specific rather than universal. Patch (applies to What the port does
One gotcha for whoever reviews it: Activation confirmed in the log rather than assumed: MeasurementRTX 3090 24 GB, Windows 11 native (not WSL), MSVC 19.44, CUDA 13.1, 96 GB DDR4, NVMe, 250 W.
Both cycles agree, SD 2–19 within each cell. Greedy output is identical between the two modes. Why I think it goes the other way hereYour DGX Spark result reads to me as "mmap is pathologically slow on that platform" rather than "explicit reads are fast" — you said as much yourself. On NTFS with 96 GB of RAM the 27 GiB table sits warm in the file cache, so an mmap fault is nearly free, and what is left is one That would predict the direct path wins on Windows exactly where the table does not fit in the file cache — a smaller-RAM box, or a machine under memory pressure. I cannot produce that state on this one (evicting 27 GiB of standby list reliably is its own project), so I have not measured it. Suggestion: if you take the patch, it may be worth leaving Two smaller notes:
|
This could help, for sure. Do you have a download link? |
You have to use a model bigger than your memory in order to test this PR or mmap. This Q3 model is too small, so there is no cache eviction. Q4 or Q5 are more interesting tests with this machine. |
…eam ggml-org#28136, adapted) The per-layer n-gram table (26.8 GiB, IQ4_NL) stays a lazy mmap, and every gathered row was a demand fault of a 4 KB page for a 90-byte row: the threaded gather plus MADV_WILLNEED took a cold 37k real-text prefill from 211 to 368 t/s on GB10 and stopped there. coder543's reader (llama-lazy- reader.h, taken as-is) serves the rows of a ubatch with explicit preads on a thread pool from an independently opened descriptor advised POSIX_FADV_RANDOM, dequantized to F32, so the table's pages are never touched; they measured 300 -> 750-800 t/s on a Spark against the plain mmap path. Adapted to our tree: the model keeps a reader per PLE tensor (load_lazy_reader), the PLE graph input carries either the row indices (mmap path) or the staged F32 rows (direct path), set_input gathers on the reader and hands the graph the same [ple_head_dim, n_heads * n_tokens] F32 tensor ggml_get_rows would have produced, so nothing downstream changes. The loader treats on-direct like on for the lazy decision. Gated bit-identical against the mmap path (tiny-model logits, and the real model's PPL 1.8175 on 10 chunks either way). Against our threaded gather it loses: cold real-text prefill with the table evicted, GB10, P70, 37k mmap+threads+willneed 361.2 t/s on-direct 330.2 t/s 140k 183.9 t/s 169.1 t/s so mmap stays the default and on-direct is opt-in (-lzm on-direct); the code is kept as upstream carries it so the next rebase takes it clean.
|
Data point from the other side of the hardware fence: on an M5 Max 128GB on Metal, on-direct is a wash to marginally slower warm. The value for me is elsewhere: under mmap, a large varied prefill fault-accumulates the PLE pages into the page cache as an ~11 GB transient burst, which on a unified-memory box eats the headroom it shares with weights and KV, tipping the machine into memory pressure. With --lazy-mode on-direct the resident footprint stays flat — a ~105k-token representative varied prefill here held wired flat with just +232 MiB pageins total and zero swap — which is exactly what lets me run a larger backbone and keep a higher-precision PLE table (Q8 engrams) that mmap would otherwise push into OOM. I'd frame this as a memory-headroom feature as much as a prefill-speed one, worth keeping even on hardware where the speedup doesn't reproduce. |
|
@eiffel31 you were right, and my earlier Windows numbers were measuring the wrong thing. Redone on What was wrong with the Q3 measurementSame box as before (RTX 3090 24 GB, Windows 11 native, MSVC 19.44, CUDA 13.1, 96 GB DDR4, NVMe, 250 W), same Windows port of the direct-read path (gist).
Cold protocol on WindowsFor anyone else measuring this on Windows: It is not sufficient on its own: mmap'd PLE pages touched by a previous request live in the process working set and survive the flush. My first attempt showed a "cold" repeat running at 515 tok/s with 89 hard faults, versus 231 tok/s and 28,325 hard faults for the genuine first run. So every cold data point below is a fresh server process plus an emptied standby list. Setup: Prefill, tok/s
The cold gain decays with prompt length (+133% → +91%) as prefill compute grows against a roughly fixed PLE fetch cost, which matches @michal-zurkowski's Strix Halo trend. Counters (cold, 29,657 tok)
Note the disk touches do not go down — both arms fetch a similar amount, and the 90 bytes per call is exactly one PLE row. What changes is that the fetches stop being serialized demand faults. That independently reproduces @nkoriyama's Linux finding (1 worker only ~1.15x, the worker pool another ~6.6x) on a completely different OS and I/O stack. Revised recommendationMy earlier note said to keep Happy to run further cells — this is a standing native-Windows install with both quants on disk. |
…eam ggml-org#28136, adapted) The per-layer n-gram table (26.8 GiB, IQ4_NL) stays a lazy mmap, and every gathered row was a demand fault of a 4 KB page for a 90-byte row: the threaded gather plus MADV_WILLNEED took a cold 37k real-text prefill from 211 to 368 t/s on GB10 and stopped there. coder543's reader (llama-lazy- reader.h, taken as-is) serves the rows of a ubatch with explicit preads on a thread pool from an independently opened descriptor advised POSIX_FADV_RANDOM, dequantized to F32, so the table's pages are never touched; they measured 300 -> 750-800 t/s on a Spark against the plain mmap path. Adapted to our tree: the model keeps a reader per PLE tensor (load_lazy_reader), the PLE graph input carries either the row indices (mmap path) or the staged F32 rows (direct path), set_input gathers on the reader and hands the graph the same [ple_head_dim, n_heads * n_tokens] F32 tensor ggml_get_rows would have produced, so nothing downstream changes. The loader treats on-direct like on for the lazy decision. Gated bit-identical against the mmap path (tiny-model logits, and the real model's PPL 1.8175 on 10 chunks either way). Against our threaded gather it loses: cold real-text prefill with the table evicted, GB10, P70, 37k mmap+threads+willneed 361.2 t/s on-direct 330.2 t/s 140k 183.9 t/s 169.1 t/s so mmap stays the default and on-direct is opt-in (-lzm on-direct); the code is kept as upstream carries it so the next rebase takes it clean.
|
I just tested this PR because I can't wait to get this improvement on Qwen Next, I figured the numbers might be worth adding here. TLDR: +167% cold prefill at 3k, decaying to +11% at 200k, and the warm cost stays under 6% My setup is a bit different from what I've seen in the thread. Linux, Blackwell, two mismatched GPUs, and not much RAM.
I have 64 GB of VRAM and 61 GB of RAM, so with a 82 GB model the PLE table can never stay cached. I think that's the normal case for people running this at home, not the exception. Same binary for both test suites (auto and on-direct), only Cold prefill across depths. I used your three lengths so it lines up with the Windows numbers, then went to deeper context:
I ran the warm cache test at every depth too, because what I'd want to know as a reviewer is whether this hurts the common case for people whose table does fit in cache:
So the warm cost is real but bounded, worst around -6% at 30k, and it basically disappears past 130k. That lines up with the 4-5% reported on Windows. On the cold prefill, I get the same decay shape as @mjungnickel18's numbers, just carried further. The gain of on-direct keeps shrinking as prefill compute grows against a roughly fixed PLE fetch cost. At 130k it's down to +22% and at 200k to +11% (268.3s to 241.4s). Still worth having! But we cannot expect the 2x prefill speed improvement at long context. On the Q3 vs Q4 correction above, I got caught by the same trap. My original benchmarks built context by repeating one sentence thousands of times and they gave me ~1818 tok/s. Real workloads on the same server gave ~609 tok/s. It was cache state. A repeated sentence keeps reading the same small part of the table, so it stays in cache and you only ever measure the warm case! Any benchmark built from repeated filler will show this PR doing nothing. Probably worth a line in the PR description, since repeating a sentence to build context is a pretty common shortcut. Let me know if there are other tests you'd want on my hardware. Thanks for this PR, it makes a real difference for local Qwen Next. |
Overview
For Qwen3.8-Flash-Next, I've been confused about the very inconsistent prefill speeds. A simple benchmark would show 700+ tok/s, so then I would start a real task, and suddenly I'm only seeing 300 tok/s. Very frustrating. This PR yields a 2x to 3x improvement in real world use, at least in my testing on my DGX Spark.
I spent a few hours this evening digging into it. Once again, the answer is
mmap. It's alwaysmmap. I really wish Nvidia would fix whatever is going on there. The simple prefill benchmark I had been running used a lot of repeated tokens, so there was very little PLE data needed, which made prompt processing fast. On real world inputs, suddenly quite few more PLE reads were needed, which caused the performance to slow way down due tommap.I haven't tested this on any other systems, but maybe these changes are actually broadly beneficial for PLE performance?
mmapeven when well-behaved is going to cause quite a bit of over-read: likely several kilobytes of wasted reads for every ~100 bytes of useful data.This PR is a very 'direct' solution to the problem I've been seeing. In an ideal world, maybe this would even be handled by something more elegant like
io_uring. But, this works, and I tried to keep the patch as small as it reasonably could be.In my testing, this boosts performance on real world input text from about 300 tok/s up to around 750 or 800 tok/s on DGX Spark, which is far better, without requiring the PLE to be pinned to RAM.
I wanted to make this new
on-directthe default behavior for GB10 owners, but I decided there wasn't an obvious way to do that which wouldn't be controversial in PR review. Maybe if other people test this PR and find that it helps on a broader range of systems, then this could become the default 'on' mode for all supported systems, with the mmap path being an alternative/fallback option.Requirements