Repository navigation
Conversation
|
Documentation preview: https://vllm--58439.org.readthedocs.build/en/58439/ |
|
Related: #58441, the graph-pool reuse race in the pinned side-stream lookup. It is the reason this backend looks up on the current stream. |
|
Haven't read the 1398 lines, so this isn't a review of the implementation. But the whole design rests on one device capability, and I have the device, so here's an independent confirmation of the gate — queried through the same driver-API path your So the opt-in gate is satisfied on the device the PR targets, on a second Spark. I checked the driver enums (88 / 100) specifically because the runtime Two adjacent attributes from the same box, in case they're useful for the doc section: That last one being 0 is worth knowing: GB10 is coherent for pageable host memory but does not advertise direct host access to managed allocations, so "unified memory" here doesn't imply the managed-allocation path is equivalent. Your approach sidesteps that entirely by mapping the files, which reads like the right call for this device rather than an incidental one. For what it's worth I hit the same wall you describe from the other side — I couldn't run the pinned backend on my Spark either while looking at #58441, for exactly the reason in your Purpose section. So the motivating constraint isn't hypothetical. |
|
Ran your new commit on a second GB10 (sm_121, CUDA 13.0), since Both halves reproduce. At Same test file with only So the new test genuinely fails on the previous head with exactly the error you named, and the One note on scope, since it cost me nothing to check and might save you a round: 24 of the 25 pass on the old source, so the regression surface here really is just the FULL-capture path. |
|
This pull request has merge conflicts that must be resolved before it can be |
|
Rebased your branch locally to see how bad the conflict is — it's small, and the result is green on a GB10, so this should be quick for you. One conflict, one file, After that, both PLE suites pass on GB10 (sm_121, CUDA 13.0), with So your One thing I checked and got wrong on the first pass, in case you try the same probe: |
be7a84f to
86f1034
Compare
|
Rebased onto main Main moved again after your rebase, so there was more to it than the
|
…heck runs Upstream dropped the function-local current_platform import from verify_model_config, which the CUDA-only check relied on (NameError on every checkpoint_mapped config). Import it where the check runs. Signed-off-by: Jürgen Schmied <juergenschmied70@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>
4a00123 to
6af74ba
Compare
|
I've been running a version of this patch locally, and through profiling discovered that this code was hard faulting the PLE pages from a CPU thread serially, leading to per pass stalling. To resolve this, I changed things so that before we actually fault the pages, we call madvise(MADV_WILLNEED) which will cause the kernel to bring pages into page cache before interaction with the cpu page tables. This reduced per step latency from 10.4ms -> 3.1ms for this phase. This was futher improved by then having the dispatch of these MADV_WILLNEED occur during the previous layer's loading. This overlap allows the next layers PLE to load concurrently wit hthe current layer's processing, reducing this PLE loading phase further from 3.1ms -> < 0.6ms on the critical path. Together these two improvements unblocked performance / minimized stalling to near zero. jschmied#4 |
|
@waych one number that might help separate the two effects. On a GB10 with the file fully page-cache-warm ( So WILLNEED on resident pages does nothing for me, which fits what it is: it queues readahead into the page cache and doesn't populate the process's page tables. The fresh-mapping minor faults are what cost 4.7 ms there, and only the touch removes them. That means your 10.4 -> 3.1 ms is probably coming from pages that weren't resident, or from the fault-around that goes with the readahead. Worth checking |
|
@hclsys My comment is about addressing latencies introduced by the use case of NOT having the data be memory resident, and instead paged off disk. The numbers I'm posting were measured via vLLM's torch profiler and represent the time spent per pass waiting for the PLE faults to finish on CPU. I'm running with very little page cache available, ~4.5GiB runtime resident for PLE. PLE is a giant hash table that is accessed with many small random reads (smaller than page size) sparsely per token. Using mmap and storing on nvme works great, as long as we dispatch the requests in concurrently, which my PR proposes. The goal is to have the kernel issue the requests to populate page cache via the MADV_WILLNEED advice (it doesn't "read-ahead", which means to read the next blocks, which in this case is wrong, the next read is random!), and then to come back and actually read them from userland CPU so that the pages get faulted and page tables established as minor faults, making them available to the GPU. The previous approach would loop and major fault serially PLE entries from the CPU, leading to slow serialized accesses to nvme with small bits of cpu in between. MADV_WILLNEED IO dispatch allows the nvme layer to dispatch multiple IO concurrently, resulting in a large reduction of major faults as many of the later pages are already in page cache by the time the page fault occurs. Performing this feed a layer early allows for a large distribution of IOs to come in before Added: The second commit, the one that currently introduces the I'm my case I'm serving the qwen3.8-flash-next PLE in fp8 as 51.2GB on disk, within 4.5GiB of usable pagecache and am now measuring a major fault rate of < 1.0 per pass with these changes. |
|
I had claude make the following artifact explaining what we saw with profiled runs. https://claude.ai/code/artifact/f070c41b-c2e8-4666-8781-40f22a3010c2 |
|
This pull request has merge conflicts that must be resolved before it can be |
Signed-off-by: Jürgen Schmied <juergenschmied70@gmail.com> # Conflicts: # vllm/models/qwen4_exp/nvidia/ngram_embedding.py
|
This pull request has merge conflicts that must be resolved before it can be |
vllm-project#56322 (sequence parallelism for Qwen3.8-Flash-Next) changed the PLE prefetch interface: start_prefetch(ngram_ids) without hidden_states, finalize_prefetch(total_tokens) in place of the pinned forward, a reduce_results constructor argument, and a lazily allocated, Optional _prefetch_buffer. - common/ngram_embedding.py: main's _finalize_prefetch (optional ETP reduce) joins through _join_prefetch_stream(), which the pageable backend overrides as a no-op. - nvidia/ngram_embedding.py: the pageable backend takes reduce_results and the new start_prefetch signature, and keeps its device in _device (its buffer is allocated in __init__); the storage choice keeps checkpoint_mapped beside main's use_reduce_scatter; the duplicate trailing __all__ is dropped. - test_ple_pageable.py: the FULL-graph test runs start_prefetch + finalize_prefetch. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Jürgen Schmied <juergenschmied70@gmail.com>
The same resolution as the base PR's merge (vllm-project#58439): vllm-project#56322 changed the PLE prefetch interface (start_prefetch(ngram_ids), finalize_prefetch(total_tokens), reduce_results, an Optional _prefetch_buffer). The pageable backend follows it; the FULL-graph test runs start_prefetch + finalize_prefetch. The readahead fill itself is unchanged. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Jürgen Schmied <juergenschmied70@gmail.com>
|
This pull request has merge conflicts that must be resolved before it can be |
Purpose
Qwen3.8-Flash-Next cannot be served on a unified-memory GPU such as DGX Spark (GB10, 128 GB shared by CPU
and GPU) with either of the storage options #54371 introduced for its 47.7 GiB FP8 PLE table:
cpu_offload, the default): a 47.7 GiB anonymous host copy of the table, in the same physicalpool as the weights and the KV cache. It cannot be cheaply dropped and re-created under memory pressure. On a
128 GB Spark, the pinned table plus weights plus KV has hard-reset the machine.
This PR adds a third, opt-in option,
--engram-config '{"checkpoint_mapped": true}', which reads the tablein place from the checkpoint's safetensors files. On GPUs that access pageable host memory through the host
page tables (
CU_DEVICE_ATTRIBUTE_PAGEABLE_MEMORY_ACCESS_USES_HOST_PAGE_TABLES), the lookup kernel reads aread-only mapping of the files. There is no table-sized device or pinned allocation, no resident duplicate of
the table, and no CPU-gather/H2D staging path; the kernel still copies the selected rows into its output
buffer. The table's pages are clean, file-backed page-cache pages that the kernel can drop and re-read, shared
by every process that maps the same files.
Design:
Qwen4ExpPLEPageableHostEmbedding, a third backend in the [Qwen4Exp] Support UVA PLE-offload and Engram tensor parallelism #54371 hierarchy. It reuses thepinned backend's buffer, ETP reduction and finalize path; only the storage and the gather differ.
safetensors data offsets are not aligned. It writes zeros for rows outside the rank's ETP range; the
existing ETP all-reduce combines ranks.
inherited side-stream lookup, greedy outputs on GB10 were not reproducible within one server start
(identical prompts, cold vs warm: 2/8 equal, |Δlogprob| up to 1.41). The observed cause on GB10: the side
stream reads the ids tensor from the CUDA-graph memory pool after the graph has released it, and later
segments overwrite it at replay. Reading a persistent copy instead made it reproducible (8/8); [Bugfix] Join side-stream work before resuming breakable CUDA graph c… #57785's capture-time sync did not
(1/8). On the current stream: 8/8, |Δlogprob| 0, and identical across fresh starts. That race ([Bug]: Qwen4Exp PinnedHost PLE prefetch reads n-gram ids from graph-pool memory that later segments reuse #58441) is
now fixed on main by [Bugfix][Qwen4Exp] Keep pinned PLE prefetch ids out of the CUDA graph pool #58489's persistent ids buffer, which this backend also uses. The lookup stays on the
current stream, which also keeps FULL CUDA-graph capture free of a second stream: this backend does not
join the pinned side stream (a GPU test captures
start_prefetch+ forward in one FULL graph). A side-streamlookup on top of [Bugfix][Qwen4Exp] Keep pinned PLE prefetch ids out of the CUDA graph pool #58489 was not measured.
30k-token prefill (≈480k rows) took 88 s with GPU faults alone, and 1.6 s when 64 CPU threads fault
the step's rows in first. The prefetch is driven from
Qwen4ExpModelState.prepare_inputs, from a CPU hashof the same inputs, and only when
checkpoint_mappedis set. It is only a hint: the lookup is correctwhether or not it has finished.
--download-dir, subfolders andthe
model.safetensors.index.jsonfilter apply exactly as for loading. Every shard that owns rows must existonce, with its exact shape and dtype; nothing is truncated.
--load-format dummymaps private anonymous zero pages.reload_weights(weights_path=...)streams those files, so it matches and remaps the new checkpoint (duringreload processing, or at the latest on the next lookup). PLE weights delivered from memory (weight sync)
differ from the files and are rejected. A rejected load stays rejected rather than falling back to the
previous mapping, and binding commits its state only after the mapping succeeded. The cost is one extra
pass over the table, on reloads only.
dp_shared_memory(mapped pages are already shared), withembedding_across_dp(the host prefetch only sees this DP rank's requests), for architectures without this backend (DeepSeek
V4.1), and off CUDA (the ROCm Qwen4Exp path from [Qwen4Exp][ROCm] PLE n-gram table CPU offload #57497 has no mapped backend and would otherwise allocate the
pinned table silently). A warning is logged on non-integrated GPUs that report the attribute, such as Grace Hopper or Grace
Blackwell: validated on DGX Spark only.
Test Plan
Unit tests (
tests/models/qwen4_exp/test_ple_pageable.py, wired into the existingmodels_basicQwen4Expjob):
dp_shared_memory,embedding_across_dp);no sample would hit; a rejected load stays rejected and a later reload from disk recovers.
start_prefetch+ forward, replayed with new ids;End-to-end on DGX Spark (GB10, sm_121, TP=1, driver 580.178, CUDA 13.0, torch 2.13, this branch's code on
the
1ea7c63f4nightly wheel). Qwen3.8-Flash-Next with an NVFP4 body, the FP8 PLE table (47.68 GiB,320,001,536 × 160 B in 128 shards / 10 files), MTP n=3.
--kv-cache-memory-bytes31 GiB unless noted.These were measured before the rebase onto current main (below); the rebased head is covered by the unit tests.
Test Result
Unit tests:
Head rebased onto main
378504a54, which includes #57497, #58489, #56926 and #58086.models_basicQwen4Exp job set (test_config.py,test_ple.py,test_ple_pageable.py): 53 passed,42 skipped (GPU-only) on CPU.
test_ple_pageable.pyon GB10: 28/28 (22 CPU + 6 GPU).tests/test_config.py -k engram: 39 passed.Memory. There is no working baseline on unmodified main (see Purpose), so the reference is the previous
offload implementation (#53899's worker, never merged) serving the same checkpoint:
checkpoint_mappedSpeed, two fresh starts per arm:
checkpoint_mappedCorrectness:
4,096-token prefill chunks.
stream-ordered async copies and no sync before the consumer. The ids equal a CPU hash of the recorded inputs,
and the rows the model consumed equal the checkpoint bytes, including all 255 mixed prefill+decode steps.
The only id mismatches were 8 CUDA-graph capture-time calls, where the ids buffer is not yet computed; those
are not serving steps.
fp8 * scalein BF16, mapped as 320 B rows)generates bit-identical output to the FP8 table: 8/8 sequential prompts and 8/8 determinism probes.
Auto KV sizing (no
--kv-cache-memory-bytes): KV sized itself to about 33 GiB. 20 min of c=8 diversetraffic (2k–60k-token prompts): 248 requests, 0 errors, MemAvailable ≥ 2.0 GiB, swap ≤ 6.8 GiB, memory PSI
full avg10≤ 6.8 % (p90 3.3 %). The speed numbers above use an explicit 31 GiB KV cap, so both arms arecompared at equal KV.
Not covered:
prefill is slow.
Comparison with other PLE-table storage implementations
What each implementation documents, plus what we measured ourselves on GB10. Other PRs' numbers are not repeated
here, and nothing was measured for #54129.
CAP_SYS_PTRACEVLLM_PLE_DISK_OFFLOAD_DIRRelation to #54129. Both keep the table file-backed and reclaimable in the page cache, and share the
discovery and validation concerns; this PR borrows #54129's approach to resolving and validating checkpoint
files.
mapping directly, with no host-side gather or staging copy.
The two could share discovery code, or
checkpoint_mappedcould become the unified-memory mode of onefile-backed backend. We are happy to rebase onto #54129 if it lands first. SGLang has an NVMe-backed path for
the same table (sgl-project/sglang#36567).
Open design questions
DefaultModelLoader._prepare_weights, a private method, to get exactly thefiles the loader read (download dir, index filter). Should the loader expose the resolved file list instead?
Its return value already changed once under this PR ([EPD][Model Loader] Skip language-model checkpoint shards for
--mm-encoder-only#58086); a test now calls the real loader.Qwen4ExpModelState. It is only a hint; thelookup is correct without it, but cold prefill is about 50× slower. Is the model state the right home, or
should it hang off a runner hook?
access, where the pinned default does not fit?
Related
unified-memory variant: the GPU reads the mapping directly, so there is no host gather and no staging copy.
The validation approach borrows from it.
and fixed on main by [Bugfix][Qwen4Exp] Keep pinned PLE prefetch ids out of the CUDA graph pool #58489.
side-stream lookup on GB10: it did not remove the non-reproducibility above (1/8).
checkpoint_mappedis rejected there.tables. They are orthogonal.
Essential Elements of an Effective PR Description Checklist
docs/features/engram.md.cc @peakcrosser7 (#54371) @Trosfy (#54129)
AI assistance: this PR was developed with Claude (Anthropic). The human submitter reviewed every changed
line, ran the tests and measurements above on the hardware named, and takes responsibility for the change.