Pull requests / #2835

#2835 prefill: the prompt path's MMQ scratch rides in the loan (opt-in, STRATA_MMQ_SLAB)

open · @gopinath87607 · 0 comentários · No GitHub

AMD / HIPNVIDIA / CUDA

Descrição

## What this changes

`STRATA_MMQ_SLAB=1` (opt-in; unset is off): the prompt path's MMQ scratch — its stream-k fixup tiles — comes out of
the loan's unused tail instead of free VRAM.

llama.cpp's `mmq.cuh` allocates those 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, and a malloc that fails there aborts the engine mid-read instead of failing the request. 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).

The pool takes the tiles from there 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.

From the Maya checkout (`project-maya`), local commit `bf844d5`.

## What it leaves alone

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.

## Measurements

- That commit 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 and no change to the chunk arithmetic, on its live split
  (5060/5060/3060/3060).
- **No speedup is claimed here.** What the start's check line shows is whether the slab is in use on a given rig.

## Checks

- **Default path:** unset = off by construction (`v != nullptr && std::atoi(v) != 0`).
- **Build:** `src/prefill/prefill.cpp` compiles clean with the build's own flags (`-fsyntax-only`);
  `src/prefill/ggml_cuda_host.cu` and `src/prefill/moe_moe.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.
- **HIP / SYCL:** not compiled here (no toolchain on the machine). The device files this touches are the CUDA MMQ
  target's; the slab plumbing in `prefill.cpp` has no per-backend branch.
- **On a start:** `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 at all, and a prompt should read
  to completion as before.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

No site

Links install, modelos, releases.