Skip to content

[Engine] the MMQ scratch now comes out of the prompt loan - #3312

Open
gopinath87607 wants to merge 1 commit into
Niko1221:mainfrom
gopinath87607:mmq-slab-loan-tail
Open

gopinath87607 wants to merge 1 commit into
Niko1221:mainfrom
gopinath87607:mmq-slab-loan-tail

Conversation

@gopinath87607

@gopinath87607 gopinath87607 commented Oct 20, 2026 •

Copy link
Copy Markdown
Contributor

Why this PR was recreated. The original (#2835) was opened from the gopinath87607/Strata fork. Deleting that fork deleted its branches, and GitHub closes every pull request whose branch lived in a deleted fork - a closed PR whose source branch is gone cannot be reopened. This is the same work from a new fork, rebased onto current main (e88ee39d); the conflict with warm_gemm_libraries beside the MMQ loan in prefill.cpp was resolved. It replaces the closed #2835, which was closed by the fork's deletion, not by a maintainer.

Problem and result

llama.cpp's mmq.cuh allocates its stream-k fixup tiles lazily through ggml_cuda_pool, which this build backs with a cudaMalloc on 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 commit bf844d5.

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 used moves the tail and the pool drops the old buffers); Prefill::run drops it once the prompt is read, before the caller refills the borrowed slots; and the loan is the participant's own expert cache cudaMalloced 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_scratch publishes nothing and mmq::pool_slab is 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 current main (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 in prefill.cpp with warm_gemm_libraries beside the MMQ loan was resolved; the branch's diff is this feature only.
  • ninja -C build strata at 78d0e621 - exit 0 (CUDA 13.2, sm_86 + sm_120).
  • -fsyntax-only on src/prefill/prefill.cpp with the build's own flags - clean.
  • src/prefill/ggml_cuda_host.cu and src/prefill/moe_mmq.cu compiled clean with the build's nvcc flags (CUDA 13.2, compute_86 and compute_120).
  • Maya checkout, local commit 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:

  • none observed.

Not run:

  • A served run with STRATA_MMQ_SLAB=1 at this head: the start's check line and a prompt reading to completion have not been observed here.
  • Serving the rebased head: only the build was re-run after the rebase; no GPU run was made at 78d0e621.
  • HIP and SYCL builds: no toolchain on this machine. The device files this touches are the CUDA MMQ target's; the slab plumbing in prefill.cpp has no per-backend branch.

Limits

  • No speedup is claimed. What the start's check line shows is whether the slab is in use on a given rig: strata prefill mmq: CUDA%d the scratch pool takes a new N MiB buffer prints for a buffer of 4 MiB or more; with STRATA_MMQ_SLAB=1 and a loan that has a tail it should not appear.
  • The slab only exists when the layout leaves a tail; a layout that used the whole loan, or spilled into owned buffers, publishes nothing and falls back to free VRAM.
  • HIP and SYCL are unverified beyond the fact that the slab plumbing has no per-backend branch.

🤖 Generated with Claude Code

@github-actions github-actions Bot added the needs-format Title or template does not follow CONTRIBUTING.md yet label Oct 20, 2026
@github-actions

Copy link
Copy Markdown

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:

  • Start the title with its area: [Engine], [GUI], [TUI], [Benchmark], [Docs] or [Other], then a short description. Example: "[Engine] fix decode stop on long prompts".
  • Fill in "Validation": the exact commands, the revision and environment tested, what you saw. Keep Passed, Failed and Not run apart.

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 /reopen.

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 gopinath87607 changed the title prefill: the prompt path's MMQ scratch rides in the loan (opt-in, STRATA_MMQ_SLAB) prefill: the MMQ scratch now comes out of the prompt loan Oct 20, 2026
@gopinath87607 gopinath87607 changed the title prefill: the MMQ scratch now comes out of the prompt loan [Engine] the MMQ scratch now comes out of the prompt loan Oct 20, 2026
@github-actions github-actions Bot removed the needs-format Title or template does not follow CONTRIBUTING.md yet label Oct 20, 2026
@github-actions

Copy link
Copy Markdown

Thanks, this follows the format now.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant