Pull requests / #3312
#3312 [Engine] the MMQ scratch now comes out of the prompt loan
open · @gopinath87607 · 0 comments · View on GitHub
Multi-GPUAMD / HIPNVIDIA / CUDALinux
Description
> **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 `cudaMalloc`ed 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](https://claude.com/claude-code)
Related on this site
Quick links to related pages.