fix: split page create from map to avoid a remap race on Volta - #465
Open
andrewleech wants to merge 2 commits into
Open
fix: split page create from map to avoid a remap race on Volta#465andrewleech wants to merge 2 commits into
andrewleech wants to merge 2 commits into
Conversation
With KVCACHED_CONTIGUOUS_LAYOUT=true and page prealloc on, the background prealloc thread does the VMM page-table edit (mem_unmap the shared zero page, then cuMemMap a real page) over a large compound page while compute kernels are in flight. On Volta this faults the running kernel and surfaces as a CUDA 700 illegal memory access. It only affects hybrid GDN/Mamba models under concurrency; KVCACHED_CONTIGUOUS_LAYOUT=false has been the workaround. Split FTensor::map into prepare() (cuMemCreate only, no mapped-VA touch) and commit() (the mem_unmap/cuMemMap VA edit, must be GPU-idle). The prealloc thread now only prepares; the VA edit runs on the main thread at the alloc_page handout, device_synchronize()'d in async mode, so cuMemCreate stays off the critical path. map()=prepare()+commit() is kept for existing callers.
Real-extension GPU test (skipped without CUDA) on the contiguous layout with prealloc on: distinct-physical-page mapping after prepare()+commit(), data integrity across free->realloc (reserved-while-mapped idempotent commit), and trim() dropping prealloc-but-uncommitted reserved pages without error.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
With
KVCACHED_CONTIGUOUS_LAYOUT=trueand page prealloc on, the background prealloc threaddoes the page-table edit (
mem_unmapthe shared zero page, thencuMemMapa real page) over abig compound page while compute kernels are still running. On Volta (V100 / sm_70) that edit
faults the running kernel and comes back as a
CUDA error: an illegal memory access(700). Itonly bites hybrid GDN/Mamba models under concurrency, and
KVCACHED_CONTIGUOUS_LAYOUT=falsehasbeen the workaround so far.
The fix splits
FTensor::mapinto two halves so the two ends can run on different threads:prepare()does just the physical allocation (cuMemCreate) and stashes the page. It touchesno mapped VA, so it's fine to run alongside live kernels.
commit()does the VA edit (mem_unmapzero page +cuMemMapreal), which has to happen whilethe GPU is idle.
The prealloc thread now only prepares. The VA edit moves to the main thread at the
alloc_pagehandout via
commit_pages(), with adevice_synchronize()in async mode. The expensivecuMemCreatestays off the critical path, so elasticity and throughput are unchanged.map()=prepare()+commit()is kept for existing callers (the worker IPC map path is untouched).commit()is a no-op if the page is already mapped (a page can sit in the reserved pool stillmapped after a free), and
unmap()just drops a prepared-but-uncommitted page with no VA edit.Testing
Validated on 2x V100-SXM2-16GB (the 1Cat sm_70 vLLM fork, OvisOCR2, TP1, contiguous layout,
prealloc on). The config that used to crash within a couple of reps (~33s) now runs 96/96 clean,
the output is numerically identical to the
CONTIGUOUS_LAYOUT=falsereference (16/16 pages), thenon-contiguous path still passes, and throughput is unchanged (88.8s vs an 84.8s no-reclaim
baseline, well within noise). Added
tests/test_create_map_split.pycovering the splitaccounting: distinct physical pages after prepare/commit, data intact across free then realloc,
and
trim()dropping uncommitted reserved pages.Also validated on stock upstream vLLM 0.26.0 (kvcached
ac9680a, this PR's base): rebuilt fromthis branch in that image and ran the same repro on a Quadro RTX 8000 (SM75), contiguous=true, TP1:
48/48 clean on the config that host has always corrupted on. So the fix holds on stock upstream and
a second GPU arch.
Caveats
A few things to flag though:
here.
into, it's a different root and this PR doesn't touch it.