Repository navigation
Reuse retained warmup mappings across compiled warmup - #154
Conversation
Retaining every freed transient segment without reusing its backing made peak preparation memory grow with each compiled-forward warmup. DeepSeek V4.1 Flash TP4 ran out of memory after 45/115 and 47/113 warmups on 96 GB Blackwell GPUs, while immediate physical release completed all 113 warmups. Keep idle transient VMM mappings in a best-fit native pool during physical retention. A later warmup receives the smallest retained mapping large enough for its allocation without creating new physical memory. The mapping remains live, its retained handle survives the complete catalog, and scope completion unmaps idle mappings before releasing every retained handle. Normal package allocations continue to create their own backing. On SM120, remapping the full extent of a retained allocation after unmap succeeds, but remapping a sub-range returns CUDA_ERROR_NOT_SUPPORTED. The pool therefore hands back whole live mappings instead of remapping prefixes. The arena records the full mapped extent when a pooled mapping serves a smaller request, and reports duplicate handle accounting as retained reference bytes. Heliopause must be rebuilt for the new allocator entry point. Validation: a DeepSeek V4.1 Flash TP4 export on four RTX PRO 6000 Blackwell GPUs (two Max-Q, two Workstation Edition) warmed 113/113 compiled forwards under retention with GPU 0 near 81 GB, then captured its 149 serving graphs and published and committed its package (85 minutes). Heliopause crate tests passed (169 passed, 3 CUDA tests ignored). The full CPU suite fails the same 38 ids as a161b17 in the same environment (4,174 passed). On GPU 3, the production-routed twelve-warmup regression and the mapped-extent inventory test passed; base package_storage.py at a161b17, with the rebuilt extension, failed that regression on the second warmup by allocating a new address while the retained mapping sat idle. The full SM120 storage suite passed 12 tests; its one mailbox-index failure reproduces on base. Material limitations: the DeepSeek package was exported to an NTFS volume and was not booted here. That export ran the revision before this one, which differed only in the failed-unmap path at scope finish, a test and an unused field, and its build also carried the mixed-SKU tuning change, a local compile-package override for the mixed card names, this host's CSF checkpoint layout, the reasoning close-token fix and the GLM-5.3 Flash TrellisMX recipe. Best-fit reuse bounds retained mappings by the distinct transient extents in the catalog, not by the largest one. SM103, SM121 and Windows were not tested. The Qwen3.8 27B DGX Spark preparation that motivated retention was not re-run under mapping reuse; reuse of retained physical memory goes beyond what a794c99 measured, though every physical reference still survives the catalog. Design-Signed-Off-By: Claude Opus 5.5 (Claude Code); design-review=1c01
finish_retention peeked the first mapping of the smallest size bucket but popped the last one. With two or more idle mappings of one size, it unmapped the first twice and removed the others from the pool still mapped, leaving arena granules that no mapping record covers. A GLM-5.3 Flash TP4 export then failed at package export on every rank with "external allocation ...+2097152 is not device memory in the serving context", at the same offset inside each rank's arena. peek now returns the mapping pop removes next. Validation: a new allocator test pools two mappings of one size, finishes retention and requires each to be unmapped once, the pool to be empty and both ranges freed once; it fails before this change (the first address is unmapped twice and the second never) and passes with it. Heliopause crate tests: 170 passed, 3 CUDA tests ignored. Design-Signed-Off-By: Claude Opus 5.5 (Claude Code); design-review=1c01
Retention finish chose the next pooled mapping with peek and removed it with pop: two expressions of one choice, which disagreed until the previous commit. Finish now takes each mapping with pop and unmaps it; if the unmap fails it puts the mapping back where it came from and returns the error, as before. peek is gone, so the drain order is stated once. The pool's invariant, that no size bucket is empty, is stated where it is maintained and asserted where pop relies on it. Validation: heliopause crate, cargo test --lib: 171 passed, 3 CUDA tests ignored. A new table test drains an empty pool, one bucket with one mapping, three buckets with one mapping each, two buckets with several mappings each, and a pool whose only small bucket take emptied; every mapping is unmapped once in drain order and its range returns to the free list once. retention_finish_unmaps_every_pooled_mapping_once and retention_finish_keeps_mapping_when_unmap_fails pass unchanged. rustfmt --check reports nothing for the file. Limitations: no behavior change is intended, so no GPU or export run was repeated for this commit; the TP4 exports ran with the previous commit's drain, which this one restates without changing its order or its failure handling. Design-Signed-Off-By: Claude Opus 5.5 (Claude Code); design-review=1c01
Retention finish removes a pooled mapping before unmapping it, and puts it back if the unmap fails, which restores the state the earlier peek-based loop kept by never removing it. A new test fails the second unmap of a two-bucket pool, and separately the first unmap of a pool whose bucket pop empties, and checks the restored pool's size, its bytes, the order a later drain removes the mappings, the free list, the retained handles and the retained reference bytes. take now asserts the same non-empty-bucket invariant pop does. Validation: heliopause crate, cargo test --lib: 172 passed, 3 CUDA tests ignored. The new test also passes on 895ba71, as does the previous commit's table test: both pin behavior this branch already had rather than showing a change. rustfmt --check reports nothing for the file. Limitations: tests and a debug assertion only; no GPU or export run was repeated. Design-Signed-Off-By: Claude Opus 5.5 (Claude Code); design-review=1c01
|
GB10 (SM121) data for the same approach, from an unpublished variant that also reuses retained transient backing across compiled warmups. Its mechanics differ from this PR (listed below), and this PR's branch itself was not run on GB10. Conditions: DeepSeek V4.1 Flash, TP4 across four DGX Spark hosts, 8192-token prefill chunk, Flare main a161b17.
How that variant differs from this PR:
Two points on this PR:
|
On the non-retain path, package_free removed the mapping record before cuMemUnmap, so an unmap error left a still-mapped range in neither the mapping table nor the free list, and a later free of the same pointer reported an unknown allocation instead of retrying the unmap. The record is now removed only after a successful unmap, as before retained-mapping reuse; the retain path still removes it before pooling the mapping. Validation: heliopause crate, cargo test --lib: 173 passed, 3 CUDA tests ignored. A new test makes the first unmap of a mapping fail and requires the record to remain and the range to stay out of the free list, then frees the same pointer again and requires it unmapped once more and freed; it fails on the previous commit, where the failed unmap already dropped the record, and passes with this change. rustfmt --check reports nothing for the file. Material limitations: the test replaces the driver's unmap with a stub that fails once; no GPU or export run was repeated for this commit. Design-Signed-Off-By: Claude Opus 5.5 (Claude Code); design-review=1c01
|
Independent reproduction on a different model and failure signature, confirming the fix class and tying it to #187. Conditions: RTX PRO 6000 Blackwell Max-Q (SM120, 96 GB), TP1 export of This PR's mechanism covers the case: pooling whole live mappings for best-fit reuse sidesteps the SM120 sub-range remap constraint ( Validation data point for the body: with a retention-reuse fix of this class, the same export completes the full 112-graph catalog at recipe defaults — peak 87 GiB with 10 GiB free through warmup — and the resulting package serves at concurrency 32 on the same 96 GB card with a 17.8 GiB KV cache (1,069,056 tokens). Decode at C=1 runs 186 tok/s; prefill ~8.6–10.2k tok/s on distinct 3.6k-token prompts. One question on the pool contract: when a pooled mapping serves a smaller allocation, the arena records the mapping's full extent — does |
Thanks for the independent reproduction, and for closing #190 in favor of this. |
|
I tested on #190 but it should be fine. |
|
Possible root cause, per #137: the CUDA_ERROR_ILLEGAL_INSTRUCTION that retention guards against may be a driver bug that cb6187b's kernel reloads expose. DGX Spark is affected because DGX OS ships driver 580.178.04, the only driver NVIDIA validates on Spark. Two Spark users in #137 report that upgrading to 615.71.09 avoids the fault with no code changes. Outside Spark, I haven't seen any reports of the original fault. Could retention apply only where the fault occurs? That could be drivers before the fixed version, or another concrete condition you prefer. It would remove the retention overhead on unaffected systems. #154's reuse would then matter only on affected Spark setups, and once Spark moves to a fixed driver, retention might not be needed at all. This would need testing on Spark first |
Reuse transient mappings only when their size equals the allocation request rounded to CUDA granularity. This keeps native mapping inventory aligned with PyTorch segments and prevents smaller persistent allocations from pinning oversized backing or exporting untracked padding. Preserve the first fatal allocator error across subsequent allocation and retention failures. Recoverable cuMemCreate OOM responses remain unrecorded so PyTorch can release cached segments and retry. Add regressions for exact-size buckets, rounded requests, error preservation after abort, and a live PyTorch allocation surviving retention completion and CUDA graph replay. Validation: - Heliopause library tests: 175 passed, 3 CUDA tests ignored. - GB10 SM121 preparation-storage CUDA suite: 12 passed, 1 mailbox-index failure that also reproduces on the unmodified PR. - Targeted retention and compiled-warmup CUDA tests: 3 passed. - The existing OOM-retry test passed in the full suite but timed out on repeat with both allocator versions. Its oversized allocation caused host OOM and flare-agent restarts on both versions. - The Qwen3.8 27B Spark crash reproduction and full-model export/memory measurements were not rerun with exact-size reuse. Design-Signed-Off-By: Codex (GPT-6); design-review=1c01
…ages Since NVIDIA#167, DeepSeek V4.1 declares a top-p range of [0.95, 1]. Vocabulary-parallel sampling supports only unfiltered target distributions: its softmax marks a row invalid when top-p is not 1, min-p is not 0, or top-k is active. A package exported with --vocabulary-parallel-sampling after NVIDIA#167 therefore admits top_p 0.95, every speculative round of that request fails verification, and the request ends with HTTP 500 "prepared model reported device error 6" (INVALID_SPECULATION). The same package serves top_p 0.95 with --mtp-draft-tokens 0, and packages exported before NVIDIA#167 clamp top-p to 1. SamplingDefaults.unfiltered() returns the policy without top-k, top-p, and min-p ranges and rejects defaults that enable those filters. DeepSeek V4.1's PreparationConfig packages that policy when vocabulary-parallel sampling is enabled and validates it on construction. As for any absent range, requests are clamped to the default (top-p 1), which keeps the documented requirement that these filters stay disabled for this option. Other ranges, including repetition_penalty [1, 1.1], are unchanged, and packages exported without the option keep the top-p range. The preparation and API documents describe the packaged policy. Validation: - New tests cover the unfiltered policy (non-filter ranges kept for top-k -1 and 0, enabled filter defaults rejected) and the DeepSeek export adapter with and without --vocabulary-parallel-sampling. The adapter test fails without the bootstrap change. - tests/preparation/test_frontend_metadata.py, tests/recipes/test_declarations.py, tests/cli/test_planning.py, tests/cli/test_speculative_graph_budget.py, tests/recipes/test_deepseek41_frontend.py, tests/preparation/test_native_recipe_contract.py and tests/preparation/test_runtime_package.py, with Heliopause built with the qualification feature: 296 passed, 37 skipped, 11 failed. The same 11 fail on cc696a6 without this change: five need the Kimi K3 checkpoint, two are the deepseek41 and muse_spark default-sampling expectations, two are runtime collective-ladder cases, and three are graph-budget cases. - GPU: exported a derivative of the qualified DeepSeek-V4.1-Flash-1017 revision at TP4 on four RTX PRO 6000 Blackwell GPUs (C32, MTP 7, --vocabulary-parallel-sampling, NVFP4 draft head, FP8 Markov head; NVIDIA#154 applied locally so the C32 export fits) and served it with flare run. flare inspect shows no top-p range and keeps the repetition-penalty range; a package exported from 902c88e shows top-p [0.95, 1]. Chat requests with top_p 0.95 at C1 and C8 and one streamed request all returned 200 with 3.4-3.8 accepted drafts per verification round, where the 902c88e package returned device error 6. Functional API checks passed 34/34 and decode throughput was unchanged. Limitations: on these packages a request for top_p below 1 receives top-p 1, as before NVIDIA#167, rather than a filtered distribution. Top-p support in vocabulary-parallel verification would change its sampling kernels and is not attempted here. Only single-host TP4 was exercised. Design-Signed-Off-By: Claude Opus 5.5 (agent for Josh Cartu); design-review=a38f Signed-off-by: Josh Cartu <j@mail99.me>
…mappings c468f0d scoped PackageStorage.retain_transient_backing() to integrated GPUs because upstream a794c99 retained every freed transient segment across the whole compiled-forward warmup and the 149-forward qwen38_dflash export ran out of memory on a 96 GB GB202. Upstream 680be69 (NVIDIA#154) now keeps idle transient mappings in a best-fit pool and reuses them, which is the fix the gate worked around. capture.py and its retention test return to upstream's text. Validation: probe export of qwen3.8-flash-dflash2-yarn4 (mixed prefill, concurrency 8, fixed-7, 149 forwards) from image 141d152f with this revert on GPU 0: Warm compiled forwards 149/149 in 453 s, entry gb202-tp1-202611041345 (173 graphs, max_active 8, retained storage 79,675 MiB at export); tests/preparation/test_transient_retention.py 3 passed on the host. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Design-Signed-Off-By: Claude Fable 5.1 (claude-fable-5-1); design-review=a38f
…mappings c468f0d scoped PackageStorage.retain_transient_backing() to integrated GPUs because upstream a794c99 retained every freed transient segment across the whole compiled-forward warmup and the 149-forward qwen38_dflash export ran out of memory on a 96 GB GB202. Upstream 680be69 (NVIDIA#154) now keeps idle transient mappings in a best-fit pool and reuses them, which is the fix the gate worked around. capture.py and its retention test return to upstream's text. Validation: probe export of qwen3.8-flash-dflash2-yarn4 (mixed prefill, concurrency 8, fixed-7, 149 forwards) from image 141d152f with this revert on GPU 0: Warm compiled forwards 149/149 in 453 s, entry gb202-tp1-202611041345 (173 graphs, max_active 8, retained storage 79,675 MiB at export); tests/preparation/test_transient_retention.py 3 passed on the host. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Design-Signed-Off-By: Claude Fable 5.1 (claude-fable-5-1); design-review=a38f
a794c99dretains the physical backing of every transient allocation freed during the compiled-forward warmupcatalog, on all supported hardware, and releases it before capture. Each warmup's transient pool creates new
physical memory, so the retained backing grows with every compiled forward. On four 96 GB RTX PRO 6000 Blackwell
cards, DeepSeek V4.1 Flash TP4 preparation ran out of memory in "Warm compiled forwards": 45 of 115 forwards with
111 MiB free on GPU 0, and 47 of 113 with 49 MiB free ("this process has 94.95 GiB memory in use"). The same
export with transient backing released immediately completed all 113 forwards.
Change
Retention is unchanged on every GPU: no physical allocation freed during the catalog is released before the
scope finishes. What changes is reuse. A transient free during retention keeps its VMM mapping live, takes its
retained physical reference as before, and puts the mapping in a best-fit native pool. A later transient
allocation in the same scope receives the smallest pooled mapping large enough for it instead of new physical
memory. Scope completion unmaps every idle pooled mapping before it releases every retained reference, and a
failed scope still makes no driver release call. Normal package and graph-pool allocations never draw from the
pool.
flare_package_transient_allocate, serves the transient pool; Heliopausemust be rebuilt for it.
returns
CUDA_ERROR_NOT_SUPPORTED, so the pool hands back whole live mappings instead of remapping prefixes.export see the real mapping. The snapshot reports duplicate references to reused memory as
retained_reference_bytes, and idle pooled mappings asretained_reusable_handlesandretained_reusable_bytes.mappings of one size, finish unmapped one of them twice, returned that range to the arena's free list twice, and
dropped the other from the pool while it was still mapped. A GLM-5.3 Flash TP4 export on this host then failed
at package export on every rank, for example
external allocation 0xb0cdfa00000+2097152 is not device memory in the serving context (memory 2, context 0x0, ...). A new allocator test fails without the change, with the firstrange in the free list twice, and passes with it.
popalone: finish removes each mapping and unmaps it, and if the unmapfails it puts the mapping back where it came from, which is the state the earlier loop kept by never removing
it. The drain order is stated once and no behavior changes; a table test drains five pool shapes and checks every
mapping is unmapped once and its range freed once.
drain follows, the free list and the retained handles) and gives
takethe same non-empty-bucket assertion aspop. The third and fourth commits' tests also pass on the second commit: they pin behavior rather than changeit.
removal ahead of
cuMemUnmapon the non-retain path, so an unmap error left a still-mapped range in neither themapping table nor the free list; the order is now the one before retained-mapping reuse, and a new test fails the
first unmap, requires the record to remain, then requires a retry to unmap and free the range.
Validation
a161b17fwith this change:
Warm compiled forwards. 113/113 171.1sunder retention, with GPU 0 at about 81.1 GB inone-minute samples through warmup (the failing runs held 94.95 GiB). The export then captured 149 serving graphs,
verified and published all four ranks' data and committed the signed package (rc 0, 85 minutes). That run used
an earlier revision of the first commit, which drained the pool with
popand had no failed-unmap handling; thefinal branch drains in the same order. The other differences are tests and an unused field.
best-fit selection, no driver call after an aborted scope, keeping a pooled mapping whose unmap fails,
unmapping every pooled mapping exactly once at finish, draining five pool shapes, the pool's exact state after a
failed unmap, and a freed mapping that keeps its record when its unmap fails.
GLM-5.3 Flash TP4 export that had failed at package export completed on this host in two later windows. No export
was repeated on the third, fourth or fifth commit.
tests/preparation/test_package_storage_cuda.py, one RTX PRO 6000 Blackwell): twelve warmups through thereal
PackageStorage.retain_transient_backing()andtransient()path reuse one mapping, read from theallocator's own snapshot; the mapped-extent inventory test passes. With base
package_storage.pyata161b17fand the rebuilt extension, the twelve-warmup test fails on the second warmup by allocating a newaddress while the retained mapping sits idle. The suite's one mailbox-index failure reproduces on the base.
failed, 4,174 passed, 2,462 skipped on unmodified
a161b17fin the same environment. The failing ids areidentical; the two extra skips are the new CUDA tests.
Limitations
was not booted there; this change does not touch serving or package contents.
override of torch's compile-package GPU-name check for the two card names, this host's CSF checkpoint layout,
the reasoning close-token output fix, and the GLM-5.3 Flash TrellisMX recipe.
one; the DeepSeek V4.1 catalog fits, and other catalogs were not measured.
re-run under mapping reuse; reusing retained physical memory goes beyond what
a794c99dmeasured, though everyphysical reference still survives the catalog.