qwen4exp: direct reads for the lazy PLE table (>2x prefill performance improvement on GB10)

[coder543]

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 always mmap. 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 to mmap.

I haven't tested this on any other systems, but maybe these changes are actually broadly beneficial for PLE performance? mmap even 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-direct the 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

• I have read and agree with the contributing guidelines Yes
• AI usage disclosure: Yes. This is heavily written and reviewed by AI. GLM-5.3 wrote the code, GPT-5.6 Sol reviewed it across multiple rounds to whittle away at rough edges. I've also reviewed the code myself and it seems fine to me now. I've tested it and does work on my system.

[nkoriyama]

I independently tested --lazy-mode on-direct on an RX 9070 XT / Linux / Vulkan / NVMe system using a real 23,664-token prompt and a strict-cold procedure: the PLE byte range is evicted before each cold run and verified with mincore; when posix_fadvise(DONTNEED) does not meet the cold-residency threshold, the server is restarted to release the mapping.

The measurements below use a build with instrumentation-only changes to record pread/dequant timing and RUSAGE_THREAD counters; the direct-read execution path itself is unchanged from the PR.

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 get_rows row loop (208,165 process-wide, 99.88%), costing 18.01 s in that serialized row loop, or about 86.6 µs per major fault.

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:

path	events / reads	cold local/stage wall
mmap demand faults	207,913 faults	18.01 s
direct, 1 worker	296,120 row reads	15.75 s
direct, PR default (32 workers on this 16-thread CPU)	296,189 row reads	2.39 s

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.

[Rhonstin]

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 UD-Q3_K_XL (84 GB, 3 shards) fully offloaded, with the ~6 GB IQ1_S PLE table served from an NVMe SSD via --lazy-mode — the table has to stay on disk here, since the host only has 15 GB of RAM. 160K context slot, ubatch 256, no speculative decoding during the benchmark, to isolate the PLE path.

I patched this PR onto my tree and compared --lazy-mode on vs --lazy-mode on-direct with a strict-cold procedure: sync && echo 3 > /proc/sys/vm/drop_caches before every server start, then a 65,715-token prompt (random filler, single user turn). "Warm" is a second, differently-seeded 65K prompt, so the server KV prefix cache is never hit — only the PLE page cache state differs between the two runs. pp is the server-reported prompt_per_second.

--lazy-mode	prefill 65K, cold	prefill 65K, warm
on	403 t/s (first pass: 374)	471 t/s
on-direct	552 / 558 t/s	593 / 575 t/s

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:

1. Applies cleanly except for a small context clash in load_arch_tensors() — my tree carries qwen4exp : add NextN/MTP draft head (--spec-type draft-mtp) for Qwen3.8-Flash-Next #27836's NextN/MTP head changes, so the local PLE block differs slightly from current master. The direct-reader block itself dropped in verbatim and worked on the first try.
2. With the MTP head enabled the absolute prefill numbers are lower (~240 t/s baseline), but the relative win holds there too.

Adopting on-direct for this server — thanks!

[ServeurpersoCom]

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.

[Rhonstin]

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 on was the worst case (374–403 t/s), while on-direct with the worker pool pulled NVMe out of the serialized-fault pattern (552–558 t/s).

[michal-zurkowski]

Strix Halo (gfx1151) on-direct gives +20–32% cold prefill

llama.cpp build 10743 54ed7d366 + this PR, ROCm 10.0, 128 GB unified memory (125 GiB usable), NVMe.

Model unsloth/Qwen3.8-Flash-Next-GGUF (qwen4exp). PLE table per_layer_token_embd.weight is
26.8 GiB in UD-Q4_K_XL (103.7 GiB total) and 50.7 GiB in UD-Q5_K_XL (147.4 GiB total larger than
the 125 GiB of RAM, so only loadable at all because of lazy PLE).

Method. Only --lazy-mode varies, stock defaults otherwise:

llama-server --model <gguf> -c 131072 -lm dio --lazy-mode {off|on|on-direct}

Per rep: drop_caches → start server → send prompt (cold) → send again (warm) → kill.
5 reps × 2 quants × 3 modes = 180 timed prefills, no failures. /completion with
n_predict:1, cache_prompt:false so only prefill is timed; numbers are the server's own
timings.prompt_per_second. (Verified cache_prompt:false really defeats slot reuse: the warm
request picks the same slot at f_sim_best = 1.000 and still re-evals all N tokens.)
P1/P2/P3 are three real-world app-spec prompts.

Cold prefill, tok/s (mean ± stdev, n=5)

prompt	tokens	Q4 on	Q4 on-direct		Q5 on	Q5 on-direct
P1	2392	250 ± 2	325 ± 1	+30%	256 ± 1	337 ± 2	+32%
P2	3767	264 ± 1	321 ± 1	+22%	269 ± 1	332 ± 1	+23%
P3	6490	259 ± 0	311 ± 0	+20%	266 ± 0	323 ± 1	+21%

Warm is a no-op -0.4% to +2.1% across all 12 cells, as expected once the rows are in page
cache. on-direct is also flat cold→warm, i.e. insensitive to cache state.

What changed (per cold prefill, mean, UD-Q5_K_XL; Q4 identical pattern)

prompt	on majflt	on-direct majflt	on-direct syscr	read MiB (both)
P1	30732	10	32274	139.9
P2	39126	0	45472	152.8
P3	61395	0	80733	239.8

read_bytes is byte-identical between the two modes: same data, different fetch. ~31–61k major
faults become 0, replaced by explicit preads. syscr is deterministic per prompt (same to the
digit across all 5 reps).

--lazy-mode off

Doesn't fail on the oversized model - it swaps. UD-Q5_K_XL loads in 159 s vs 19 s for
on-direct, parking ~43 GiB in swap (avail. drops to ~5 GiB), and its cold prefill is still
slower (312–318 vs 323–337). Q4: 34 s load, ~7 GiB swapped. No case for off on either quant here.

The gain decays with prompt length (+31% at 2.4k tok → +21% at 6.5k), since prefill compute grows
while PLE fetch is ~fixed per unique row.

[Farenheith]

If I understood correctly, this PR's new --lazy-mode on-direct seems to be the ideal approach for local inference of Qwen3.8-flash-next and any future models using PLE/n-gram tables.

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 -ngl 99 --lazy-mode on-direct allows offloading 100% of the weights to the accelerator while keeping the n-gram table permanently on the SSD with minimal I/O overhead.

This is gold!

[pwilkin]

I feel like this is a valid investment for the big gains it provides, although I'm not sure if Windows support should be a part of this or a followup PR. But we'd probably need @ngxson to chime in here.

[ggerganov]

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.

[gabrielfreire]

I had to add the following line to llama-mmap.h in order to compile this `

#include <string>

am I the only one who had this issue? I am on windows

Also

load_arch_tensors: --lazy-mode on-direct is not supported on this platform, using lazy mmap reads

Is windows support coming?

[eiffel31]

@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?

[coder543]

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 string there. If someone wants to contribute full Windows support, that would be great, but as mentioned, I don't have Windows available at the moment to test 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 qwen4exp does with its PLE layout, but I agree it would make sense to support as an option. When I tested Gemma 4 E2B with on-direct, I saw somewhere between "no improvement" and a 2% regression in prefill speeds with on-direct versus on.

For consistency, it would be a nice option to have on Gemma, but the benefit mostly seems to be exclusive to qwen4exp for now. I would be happy to refactor this to be shared by Gemma 4 E2B/E4B and Qwen3.8-Flash-Next if that would be helpful?

@eiffel31 the model that I was testing is this one. The command line is nothing special, what matters is setting --load-mode none and changing --lazy-mode between on and on-direct, as well as using realistic input text, not artificial prompt text that is a small string repeated thousands of times. Real text causes a lot more PLE entries to be accessed, whereas a single repeated string does not. I also used -ub 2048.