Repository navigation
[Engine] the MMQ scratch now comes out of the prompt loan - #3312
Open
gopinath87607 wants to merge 1 commit into
Open
gopinath87607 wants to merge 1 commit into
gopinath87607 wants to merge 1 commit into
Conversation
|
Thanks for the pull request. To keep the tracker readable, every item needs an area prefix in the title and the parts of the pull request template filled in. This one is missing:
Edit the pull request (title and description) and this note goes away on its own. If it is still missing parts in 5 days it will be closed, which is not a rejection: fix it and comment Details: CONTRIBUTING.md |
…ATA_MMQ_SLAB) llama.cpp's mmq.cuh allocates its stream-k fixup tiles through ggml_cuda_pool, which this build backs with a cudaMalloc on whatever VRAM is free at that moment - a few MB on a card the expert cache has filled, at the one moment a prompt is running, and a malloc that fails there aborts the engine mid-read instead of failing a request. The Maya checkout (`project-maya`) hit exactly that at a 16,384-token chunk (3.4 MB free of 7,707.7, 1.2 MB wanted) and fixed it by lending the pool the loan's unused tail (local commit `bf844d5`). A layout does not use the whole loan: the carve is a bump, so everything from its stop to the loan's end is provably the prompt's own memory (a loan is whole cache slots and the top one is only partly needed). The pool now takes the tiles from there. It is published per layout, keyed on the tail's ADDRESS, republished at the start of every run and withdrawn when the run ends - the caller refills the borrowed slots as soon as the prompt returns, and a cached buffer in a refilled slot would be silent corruption rather than a crash. A layout that uses its whole loan, or that spilled into owned buffers, publishes no tail and the pool allocates as before. The pool also rounds a request up to a size class (power of two, 128 KiB floor). Keyed on the exact size it held one buffer per J of MMQ's J search - 32 buffers, 31.5 MiB on a 579-token chunk (the Maya checkout's measurement on its rig) - where a class bounds it to a few MB, and the cudaMalloc that fails now names what the pool held, how much was free and how much the request wanted. STRATA_MMQ_SLAB=1 turns it on; unset, empty or 0 leaves the pool on free VRAM exactly as before, so the default path is unchanged and a new feature is opt-in. The size classes stay either way - they are what bounds the pool. The tokens are unchanged in both cases: the tiles are scratch and the kernels touch only the first nsm*J*I floats of whatever they are handed. Checked: `src/prefill/prefill.cpp` compiles clean with the build's own flags (`-fsyntax-only`); `src/prefill/ggml_cuda_host.cu` and `src/prefill/moe_mmq.cu` compile clean with the build's nvcc flags (CUDA 13.2, compute_86 and compute_120); the development branch that also carries the change built fully with `ninja -C build strata`. `engine/strata` was not touched. No HIP toolchain exists on this box, so the HIP compile of the three touched sources was not run - the new code uses only cudaMalloc, cudaMemGetInfo, cudaFree, cudaGetErrorString, cudaError_t and std::remove_if, all of which the HIP compat header maps or the host compiler provides. The check to run here (a prompt reads to completion, and no `strata prefill mmq: CUDA%d the scratch pool takes a new N MiB buffer` line at the start) is the owner's to run. Co-Authored-By: Claude Code <noreply@anthropic.com>
gopinath87607
force-pushed
the
mmq-slab-loan-tail
branch
from
October 20, 2026 01:57
a177944 to
78d0e62
Compare
|
Thanks, this follows the format now. |
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.
Problem and result
llama.cpp's
mmq.cuhallocates its stream-k fixup tiles lazily throughggml_cuda_pool, which this build backs with acudaMallocon whatever is free at that moment - a few MB on a card the expert cache has filled, at the one moment a prompt is running. A malloc that fails there aborts the engine mid-read instead of failing the request. After this change,STRATA_MMQ_SLAB=1(opt-in) takes those tiles from the loan's unused tail instead of free VRAM. A layout does not use the whole loan: the carve is a bump, so everything from its stop to the loan's end is the prompt's own memory, provably (a loan is whole cache slots and the top one is only partly needed).Related discussion
The original pull request and its discussion: #2835 (closed when its fork was deleted; this PR replaces it). The work comes from the Maya checkout (
project-maya), local commitbf844d5.Approach
The pool takes the tiles from the loan's tail in size classes (a power of two, 128 KiB floor). The classes are load-bearing: MMQ's J search picks a different J per launch, so a cache keyed on the exact size holds one buffer per J - 32 buffers, 31.5 MiB on a 579-token chunk in the Maya checkout's measurement - where a class bounds the pool to a few MB.
Three guards make it safe, all load-bearing: it is published per layout keyed on the tail's address (a request whose chunk moves
usedmoves the tail and the pool drops the old buffers);Prefill::rundrops it once the prompt is read, before the caller refills the borrowed slots; and the loan is the participant's own expert cachecudaMalloced on this device, so the tail is this card's memory to write. A tail that is empty (the layout used the whole loan, or spilled into owned buffers) publishes nothing, and the pool allocates as it always did.Unset (the default), the pool stays on free VRAM exactly as before:
lend_mmq_scratchpublishes nothing andmmq::pool_slabis not given a slab. The size classes are part of the slab path only (they bound the slab's pool); with the knob off there is no slab to bound. Default path unchanged by construction.Validation
Tested revision:
78d0e621, one commit on currentmain(e88ee39d). Environment: Linux, CUDA 13.2 (V13.2.86),sm_86+sm_120; the build machine is the Q8_0 4-way rig (2x RTX 3060 12 GB, 2x RTX 5060 8 GB).Passed:
git rebase e88ee39d- the conflict inprefill.cppwithwarm_gemm_librariesbeside the MMQ loan was resolved; the branch's diff is this feature only.ninja -C build strataat78d0e621- exit 0 (CUDA 13.2,sm_86+sm_120).-fsyntax-onlyonsrc/prefill/prefill.cppwith the build's own flags - clean.src/prefill/ggml_cuda_host.cuandsrc/prefill/moe_mmq.cucompiled clean with the build's nvcc flags (CUDA 13.2,compute_86andcompute_120).bf844d5(not this engine, not run by us here): it hit the failure at a 16,384-token chunk on a 5060 (3.4 MB free of 7,707.7 MB, 1.2 MB wanted) and measured the slab costing 42-47 more lent experts per part with no change to the chunk arithmetic, on its live split (5060/5060/3060/3060).Failed:
Not run:
STRATA_MMQ_SLAB=1at this head: the start's check line and a prompt reading to completion have not been observed here.78d0e621.prefill.cpphas no per-backend branch.Limits
strata prefill mmq: CUDA%d the scratch pool takes a new N MiB bufferprints for a buffer of 4 MiB or more; withSTRATA_MMQ_SLAB=1and a loan that has a tail it should not appear.🤖 Generated with Claude Code