Think managed kv - #669
Open
ruota wants to merge 3 commits into
Open
Conversation
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.
Allow managed KV override for CUDA
Summary
Add the documented
DS4_CUDA_FORCE_MANAGED_KV=1opt-in for CUDA. It forcesonly the long-lived KV allocation into CUDA managed memory, while model and
optional DSpark support weights remain resident in VRAM.
This makes it possible to run long-context experiments where the weights fit
in VRAM but the KV allocation would otherwise prevent session creation.
The change also includes
<float.h>so the CUDA backend has the standardFLT_MAXdefinition it already uses.Scope
DS4_CUDA_FORCE_MANAGED_KV=1enables the override.loading behavior.
Testing
Builds completed on an NVIDIA RTX PRO 6000 Blackwell (
sm_120):make cpumake cuda CUDA_ARCH=sm_120make cuda-regression CUDA_ARCH=sm_120git diff --checkRuntime validation used DeepSeek V4 Flash IQ2 imatrix (80.76 GiB) plus the
DSpark support GGUF (5.58 GiB) at
ctx=393216:DS4_CUDA_Q8_F16_CACHE_MB=0 \ DS4_CUDA_FORCE_MANAGED_KV=1 \ ./ds4 --cuda -m ./ds4flash.gguf \ --mtp ./gguf/DeepSeek-V4-Flash-DSpark-support.gguf \ --dspark --think-max --ctx 393216 --prefill-chunk 256 \ --tokens 32 --temp 0 -p "..."The run loaded the main model and DSpark support resident in VRAM, placed the
5.08 GiB KV allocation in managed memory, and completed at 415.66 prefill t/s
and 59.46 generation t/s in the CLI benchmark. A warm server request completed
64 tokens in 1.2667 seconds (about 50.5 end-to-end tokens/s).
The DSpark acceptance fixture was also run with the existing local GGUFs and a
256-token prefill chunk so the resident main and support models fit on the 96 GB
GPU. Managed-KV and device-KV runs produced matching target/DSpark output and
the same proposal statistics. Four cases accepted every proposed draft; the
existing
c_addproposal-quality guard failed in both modes because confidence0.9 produced no draft for that case. This is therefore not caused by the
managed-KV override.
The stock
dspark-verify-depthfixture uses a fixed 4096-token prefill bufferand cannot create its sessions with both resident GGUFs on this 96 GB GPU.
Notes
This is an explicit capacity override. Managed memory can reduce throughput if
the KV pages are not resident, so it is not enabled automatically.