feat(qwen4): stream PLE embeddings from NVMe

[jzinno]

Stack

This PR is stacked on #36497 and targets its qwen4-main-squashed branch. It should be rebased onto main after that PR lands.

Motivation

The Qwen3.8 Flash Next NVFP4 checkpoint contains a 47.68 GiB FP8 PLE n-gram embedding table. On a 128 GiB unified-memory system such as DGX Spark, keeping that table resident prevents the rest of the model, Mamba state, and KV cache from fitting comfortably.

PLE selects a small number of rows per forward, so the full table does not need to be resident. This change leaves the table in its original sharded safetensors files and reads only the selected rows from local NVMe.

Modifications

• Add a bundled Linux Rust extension with a persistent io_uring, page-aligned storage, bounded submission batches, GIL-free reads, and preserved OS error codes.
• Parse and validate sharded PLE safetensors metadata without loading the table.
• Add io_uring and mmap row readers, optional page caching, pinned staging memory, asynchronous H2D conversion, and overlap with the decoder layer before PLE.
• Skip loading resident PLE shard weights when NVMe streaming is enabled. The initial path is intentionally TP1 and FP8 E4M3 only.
• Add a correctness-first Triton sparse GQA decode fallback for SM121, used only when the preferred TRT-LLM sparse path is unavailable.
• Add native, CPU, and GPU unit tests plus a deployment guide and a crate-scoped Rust CI job.

The feature is opt-in through SGLANG_QWEN4_PLE_NVME_PATH; existing model loading is unchanged when it is unset.

Accuracy Tests

• The manifest and mmap tests cover sharded row mapping and malformed shard layouts.
• The Python binding test reads multiple pages through a queue-depth-one ring.
• Against RadixArk/Qwen3.8-Flash-Next-NVFP4 revision 7b719225242aacd3dbd3f9407468c2ee9a9d2594, rows at the first shard, a shard boundary, and the final shard matched safetensors.safe_open byte for byte through the Rust reader. Sixteen additional random rows also matched.
• The SM121 sparse GQA kernel matches a float32 grouped-query reference for the model shape: 24 Q heads, 2 KV heads, head dimension 256, and 2,051 selected positions. Maximum absolute error was 0.0009765625 and mean absolute error was 0.0000765.
• The exact branch served the full checkpoint through the Rust reader and completed a deterministic OpenAI-compatible chat-completion smoke test.

Completed locally:

cargo test -p sglang-storage                                           1 passed
cargo clippy --workspace -- -D warnings                                passed
pytest test_qwen4_ple_nvme.py test_io_uring_reader.py                   4 passed
pytest test_qsa_decode_kernel.py                                        1 passed
pre-commit on all changed files                                         passed
mint validate                                                           passed
mint broken-links --check-anchors --check-redirects                     passed

Speed Tests and Profiling

Hardware: NVIDIA DGX Spark GB10, SM121, internal Samsung NVMe, TP1, concurrency one.

The Rust reader measured 0.208 ms p50, 0.627 ms p95, and 0.944 ms p99 for 16 random logical PLE rows over 1,000 iterations while another storage-heavy workload was active. The sparse GQA kernel measured 0.0978 ms mean over 100 iterations at batch size one.

The exact branch was then run end to end with RadixArk/Qwen3.8-Flash-Next-NVFP4, BF16 KV cache, a 32K context, eager execution, and the checkpoint's built-in NEXTN head using three speculative steps, top-k one, and four draft tokens. Each domain contains ten held-out prompts with 512 output tokens per request:

Domain	Input tokens	Output tokens	Duration	Output tok/s	Accept length	Median TTFT	Median TPOT
Chat	8,162	5,120	217.72 s	23.52	2.51	939.78 ms	39.88 ms
STEM	2,892	5,120	215.26 s	23.78	2.54	462.02 ms	38.11 ms
Math	1,379	5,120	189.82 s	26.97	2.69	310.90 ms	36.76 ms
Code	6,456	5,120	222.49 s	23.01	2.68	801.03 ms	41.48 ms
Overall	18,889	20,480	845.30 s	24.23	2.60	469.02 ms	39.06 ms

All 40 requests completed without an API error, restart, or OOM. Overall TTFT is the median across all requests; overall TPOT is the mean of the four domain medians. After 7,000 logged gathers, the Rust reader had selected 703,568 rows at 4.475 ms mean read time across mixed prefill and speculative verification batches.

The target weights loaded in 466.10 seconds and occupied 80.06 GiB. The integrated NEXTN pass loaded in another 90.20 seconds. Target, draft, Mamba, and KV pools left 18.31 GiB available, with 447,040 KV-cache tokens allocated.
\n

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.
• 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 #33088727474
Latest PR Test (Extra): ❌ Run #33088725886
Latest PR Test (AMD ROCm 10): ❌ Run #33088726382

[benthecarman]

I ended up building the same thing but can confirm this strategy works and I was able to reproduce similar results.

One finding worth folding in: the docs position mmap as correctness-only, but it’s production-viable (~150 µs per decode-step gather, 16 ms per 65k-row prefill gather warm) except for an RSS problem on modern kernels. With large page-cache folios (Linux 6.x), each random row fault maps a whole folio, so process RSS silently climbs toward the full 47.7 GiB table; on unified-memory boxes that skews the free-memory reads used for KV sizing. MADV_RANDOM doesn’t help, it limits readahead I/O, not mapping-in of cached folios. The fix I run in my implementation is ~15 lines: periodically sum the Rss of just the table mappings from /proc/self/smaps and madvise(MADV_DONTNEED) them above a budget; pages stay in cache so hot rows re-fault at minor-fault cost. Happy to share the snippet if wanted

[hashd1ve]

@jzinno — we built the same feature independently, three days apart, on the same base branch, and neither PR references the other: #36567 (2026-08-26) and #37068 (2026-08-29) both take the 47.7 GiB Qwen4-Exp PLE table out of resident memory and serve its rows from NVMe. Flagging it before a reviewer has to discover it, and proposing where they converge.

They also collide mechanically: both edit python/sglang/srt/models/qwen4_exp.py (+67/-23 here, +29/-10 there), so whichever lands second needs a rebase regardless of the outcome.

Where they differ

	#36567	#37068
Mechanism	reads rows from the original sharded safetensors at inference, through a bundled Rust io_uring reader or an mmap reader, with pinned staging and async H2D	keeps the existing --ple-offload-embedding host table and moves it from pinned memory to a sparse file mapped with torch.from_file(shared=True); the existing Triton gather kernel dereferences that mapping directly
Per-gather copies	staging buffer → H2D	none: no reader, no staging, no H2D
Hardware scope	any box with local NVMe	unified memory only; validated at load time against cudaDevAttrPageableMemoryAccessUsesHostPageTables == 1
Surface	new model path qwen4_ple_nvme.py (+626), new Rust crate sglang-storage, env var SGLANG_QWEN4_PLE_NVME_PATH, +1555/-94 in 16 files	a second backend on the flag that already exists: --ple-offload-backend {pinned,file} (default pinned, behaviour unchanged), one Triton-free module, +647/-12 in 10 files
Sharding	intentionally TP1 and fp8 e4m3 to start	file name encodes shape, dtype, size and the rank's vocabulary range, so TP shards never share a file
Weight loading	skips loading the resident PLE shards	unchanged: copy_ into the mapped tensor writes through to the file, so every boot rewrites the table and a stale file cannot leak old rows

Neither is a superset. Yours is portable and doesn't depend on the GPU being able to walk host page tables; it pays a reader, a staging buffer and a copy for that. Mine is zero-copy and small enough to be a flag value rather than a parallel model path, but it only works on coherent unified-memory parts, which today means GB10.

A way to converge

They compose better than they compete: the file backend and an io_uring backend are two values of the same knob. Concretely — --ple-offload-backend {pinned,file,nvme}, with your reader behind nvme and the mapping behind file, both selected at allocation time in one module, one user-facing flag instead of a flag plus an env var plus a second model class. That keeps your portability for discrete GPUs and the zero-copy path where the device attribute allows it, and reviewers get one surface to reason about. Happy to do that merge work in either direction — land yours first and I'll rebase #37068 into it as a backend, or the reverse, whichever a maintainer prefers.

One data point for @benthecarman's RSS finding

Confirming it independently on a GB10, and with a correction that cost me time: the RSS climb is real, MADV_RANDOM does not stop it, and posix_fadvise(DONTNEED) does not release the cached folios either — it silently does nothing here. process_madvise(MADV_DONTNEED) over the table's mappings is what actually drops them, and it takes about 3.5 s for the full table on this box. Growth measured while serving is roughly 45 KB per generated token, so on a 121.6 GiB unified-memory box this is not a slow leak, it reaches the free-memory readings that size the KV pool within a single long session.

This applies to my #37068 as much as to your mmap reader — the PR as it stands advises MADV_RANDOM and prefetches with posix_fadvise(WILLNEED), and has no release path — so I'll add the periodic release regardless of how the two PRs resolve.

Unrelated, but it saves you a rebase

#36845 landed on qwen4-main-squashed today (78c5024e) with a Triton SM121 QSA decode kernel, routed unconditionally on sm_121. That makes the SM121 sparse-decode fallback in this PR (kernels/ops/attention/qsa_decode.py, and the qwen_sparse_attn_backend.py +30/-60 hunk) redundant, and it's ~180 lines of the diff plus the part most likely to conflict.