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!