[Qwen4-Exp] File-backed PLE table backend for unified-memory devices (GB10 / DGX Spark)

[hashd1ve]

Motivation

--ple-offload-embedding keeps the Qwen4-Exp PLE n-gram table (47.7 GiB in fp8 for Qwen3.8-Flash-Next) in CPU pinned memory and lets the Triton gather kernel read rows from the host pointer. On a discrete GPU that frees VRAM. On unified-memory parts — GB10 / DGX Spark today — pinned host memory comes out of the same pool as the model weights, so it frees nothing: the RadixArk/Qwen3.8-Flash-Next-NVFP4 checkpoint is 126.0 GiB of weights on a 121.63 GiB box and does not boot.

The GB10 reports cudaDevAttrPageableMemoryAccessUsesHostPageTables = 1: the GPU resolves pageable host addresses through the host page tables, so a kernel can dereference a pointer into a memory-mapped file. Backing the table with a sparse file on NVMe instead of pinned RAM makes the table's residency a page-cache matter rather than a hard reservation, and the model fits with room for a 262k context. This has been serving on one DGX Spark since 2026-08-26 as a monkeypatch (recipe); this PR is the proper backend.

Stacked on qwen4-main-squashed because the Qwen4-Exp code is not on main yet.

Modifications

• --ple-offload-backend {pinned,file} (default pinned, behaviour unchanged) and --ple-offload-dir (default $SGLANG_CACHE_DIR/ple/<model path>, one directory per checkpoint). The file backend is validated against --ple-offload-embedding, and at load time against the device attribute above (SGLANG_QWEN4_PLE_FILE_SKIP_DEVICE_CHECK=1 bypasses the check for devices that cannot be queried).
• New Triton-free module sglang/srt/models/qwen4_exp_ple_table.py:
  • allocate_ple_host_table(shape, dtype, backend, table_dir): pinned as before, or a sparse file with a deterministic name (ple_table_<dims>_<dtype>_<bytes>B.bin, reused across restarts; a file of the wrong size is recreated), mapped with torch.from_file(shared=True) and advised MADV_RANDOM (the table is pure random access, 16 rows of 160 B per token; without it the kernel readahead pulled ~560x the bytes a gather uses).
  • PleFilePrefetcher: for gathers of ≥ 2048 rows (prefill-sized; decode gathers are 16–64 rows) it computes the distinct 4 KiB pages of the requested rows and issues posix_fadvise(WILLNEED) on a background thread before the kernel launches, so the page faults are served concurrently instead of one at a time. Skipped during CUDA-graph capture. SGLANG_QWEN4_PLE_FILE_PREFETCH=0 disables it.
  • PleFileRssTrimmer: a row fault maps in a whole page-cache folio, so with large folios (Linux 6.x) the mapping's resident set climbs towards the full 47.7 GiB while a token only reads a few KB of it (measured ~45 KB of Rss growth per generated token on a GB10). On a unified-memory part that is not a slow leak, because the free-memory readings that size the KV pool come from the same pool. MADV_RANDOM does not prevent it — it bounds readahead I/O, not the mapping in of folios already in cache — and posix_fadvise(DONTNEED) does not release them; MADV_DONTNEED over the mapping does, dropping the page-table entries while the pages stay in the page cache, so hot rows come back at minor-fault cost. The trimmer reads the Rss of the table's VMAs from /proc/self/smaps and, once over SGLANG_QWEN4_PLE_FILE_RSS_BUDGET_GB (default 8 GiB, 0 disables), drops them in 1 GiB slices: one madvise over the whole table holds mmap_lock for ~3.5 s, which would stall every fault in the process including the gather kernel's. It runs on its own daemon thread rather than as a hook in the gather, because decode replays a CUDA graph and executes no Python — a hook would never fire in the phase that grows the mapping. Dropping entries under a running gather is the state this backend already handles: the file starts out entirely unfaulted and every cold row is faulted in from inside the kernel through the same host page tables. Absent where the resident set cannot be read and for the pinned backend.
• Qwen4ExpPinnedHostEmbedding takes backend/table_dir, allocates through the module, and gather() calls the prefetcher when present. The gather kernel, the prefetch stream and the CUDA graphs are untouched: they keep receiving a host pointer. The weight loader is unchanged too — copy_ into the mapped tensor writes through to the file.
• Qwen4ExpConfig carries ple_offload_backend / ple_offload_dir; load_model_utils propagates them like ple_offload_embedding.
• The file name encodes shape, dtype, size and the rank's vocabulary range (rows<start>-<end>), so tensor-parallel shards of the same shape never share a file; the default directory is per checkpoint ($SGLANG_CACHE_DIR/ple/<model path>). Every boot rewrites the whole table through the unchanged weight loader, so a stale file cannot leak old rows.
• Tests: test/registered/unit/models/test_qwen4_exp_ple_table.py (mirrors the module path) — allocator (sparse and exactly sized, writes persist and the file is reused, wrong-sized file replaced, per-rank tag, per-checkpoint default dir, unknown backend rejected, pinned path unchanged), prefetcher (page set covers row start and end, dedup, size floor, advised offsets), the /proc/self/smaps parser (sums every VMA of the table, counts a partially overlapping one, ignores unrelated mappings, reports an unreadable smaps as unknown), the trimmer against a live mapping (measures its own mapping only, drops its pages once over budget without losing what was written through them, no-op under budget, thread starts and stops), and a device test that runs the production gather kernel over a file-backed table and compares with a torch gather (skipped unless the device reports the attribute).
• Docs: rows for --ple-offload-embedding (previously undocumented), --ple-offload-backend and --ple-offload-dir in server_arguments.mdx, and the five SGLANG_QWEN4_PLE_FILE_* variables in environment_variables.mdx.
• GPU coverage lives next to the existing class tests: test/registered/kernels/ops/embeddings/test_qwen4_ple_offload.py gains two file-backend cases (bf16 parity with pinned at dims 7 and 160 including a prefill-sized gather through the page-cache hint, and an fp8 table). Drive-by: that file's source stub lacked the per-tensor weight_scale buffer the class has required since 73a2552, so it failed at construction on this branch; the stub now carries it.

Not included: a chunked initialize_dummy_weights (the fp16 staging copy of a 47.7 GiB fp8 table OOMs --load-format dummy; separate PR), and a cookbook cell (the Qwen3.8-Flash-Next cookbook page is on main, not on this branch).

One deliberate tensor.cpu(): the prefetcher syncs once per prefill-sized gather (≥ 2048 rows) to compute the page set on the host; decode-sized gathers and CUDA-graph capture never reach it.

Accuracy Tests

• test/registered/unit/models/test_qwen4_exp_ple_table.py: 20 passed, and test/registered/kernels/ops/embeddings/test_qwen4_ple_offload.py: 13 passed (10 existing + 3 file-backend) — 33 with nothing skipped, on a DGX Spark (GB10, sm_121, kernel 6.17, ext4 on NVMe, CUDA 13.0, torch 2.13.0). Pre-commit: all hooks pass.
• The trim mechanism separately, on the same box over a 256 MiB shared mapping on NVMe: resident set 256.0 -> 0.0 MiB, every page written through the mapping still correct afterwards, and re-reading five pages brings back 0.2 MiB rather than the mapping — which is the property the budget relies on.
• The same mechanism (monkeypatch on the day-0 image, identical allocation and prefetch) in production on one DGX Spark, TP=1, 262,144 context, NVFP4 routed experts + FP8 dense path, NEXTN 3/1/4:
  • GSM8K (n=200): 192/200 = 96.0 % on the NVFP4 checkpoint and 193/200 = 96.5 % with the FP8 dense path, against RadixArk's published 97.27 % (BF16 band 97.12–97.50); within noise at this n. A wrongly served table does not score here.
  • Needle-in-a-haystack with the file-backed table: 4/4 exact at each of 120k, 190k and 210k prompt tokens (with the sm_121 Triton sparse-decode fallback from fix(qsa): restore SM121 correctness with Humanize and Kernel Design Agent #36845).
  • 13-case executable code suite: 13/13.

Speed Tests and Profiling

Measured on one DGX Spark (GB10, unified LPDDR5X ~273 GB/s, NVMe), fp8 table, rows of 160 B, ple_layer_ids=[2]:

gather	cold (page cache dropped)	warm
decode, 16 rows	3.58 ms	0.12 ms
prefill, 65,536 rows	3,865 ms	6.9 ms

• Disk traffic per generated token in warm decode: 138 KB with default readahead, ~64 KB with MADV_RANDOM (2.5 KB useful; the rest is 4 KiB page granularity).
• Cold prefill: 650–750 tok/s without the WILLNEED hint, 1,000–2,100 tok/s with it (warm: 2,170–2,556 tok/s), because page faults inside the kernel are serialized while the block layer can serve ~40k IOPS with queue depth.
• Decode with a cold table costs ~17 % (code: 39.0 vs 46.7 tok/s; prose: 25.0 vs 27.3) until the working set warms; a decode-time prefetch is not possible from this layer (decode runs inside a CUDA graph and the draft tokens are produced on device), so this is documented rather than fixed.
• Startup: the sparse file persists between runs, so the ~2.5 min the first boot spends writing the table are not paid again; the rest of the load is unchanged.
• End-to-end on this box, unrelated to the table but for context: 55 tok/s single-stream on code with the FP8 dense path and an FP8 lm_head, 262k context served.

Nothing changes for pinned (default): the only new code on that path is the backend dispatch in the allocator.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations (docstrings and --help; the Qwen3.8-Flash-Next cookbook page lives on main, not on this branch — happy to add a GB10 cell there once the model lands).
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ❌ Run #33365912972
Latest PR Test (Extra): ❌ Run #33365912862
Latest PR Test (AMD ROCm 7.2): ❌ Run #33365912958

[hashd1ve]

Notes for reviewers / CI:

• This PR is stacked on qwen4-main-squashed like fix(qsa): support SM120/SM121 sparse decode paths #36556, [Qwen3.8] Fix FP8 KV cache support in QSA #36644 and fix(qsa): restore SM121 correctness with Humanize and Kernel Design Agent #36845, because the Qwen4-Exp code is not on main yet. check-changes therefore fails in check-maintenance ("rebase onto the latest main"), which is not possible for a stacked branch; I will restack onto main as soon as Introduce Qwen 3.8 Flash Next #36497 lands. Until then it would need run-ci (and bypass-maintenance) from a maintainer to run — happy to do whatever is preferred.
• Drive-by: test/registered/kernels/ops/embeddings/test_qwen4_ple_offload.py on this branch fails at construction ('SimpleNamespace' object has no attribute 'weight_scale') because the source stub predates the per-tensor weight_scale buffer the class has required since 73a2552; the stub now carries it. The same file gains two file-backend cases (bf16 parity with pinned, and an fp8 table). Ran on a GB10: 13/13 in that file, 10/10 in the new unit file.

[hashd1ve]

Cross-referencing for reviewers: #36567 (@jzinno, opened three days before this one, same qwen4-main-squashed base) solves the same problem — the 47.7 GiB PLE table served from NVMe instead of resident memory — by a different route: rows read from the original sharded safetensors through a bundled Rust io_uring reader with pinned staging and async H2D, versus the zero-copy mapped file here. Neither PR knew about the other; details and a proposal to fold them into one --ple-offload-backend knob are in #36567 (comment).

Both PRs edit python/sglang/srt/models/qwen4_exp.py, so they will conflict whichever order they land in. I'm happy to do the merge work in either direction.

[yhyang201]

These are global IDs, but the file contains only this rank's shard. Please filter to the local range and subtract the shard start before prefetching.