Pull requests / #3307
#3307 [Engine] a long prompt can borrow the MTP drafter's memory while it reads
open · @gopinath87607 · 0 commentaires · Sur GitHub
BenchmarksServer & APIMulti-GPUAMD / HIPNVIDIA / CUDAModels & quantsWindowsLinux
Description
> **Why this PR was recreated.** The original (#2830) 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, based on `477ac13c`, and GitHub reports that it merges cleanly with current `main` (`e88ee39d`). It replaces the closed #2830, which was closed by the fork's deletion, not by a maintainer. ## Problem and result A long prompt's ring is bounded by the VRAM the prompt path can borrow, so on a full card the read waits on the ring. The MTP drafter's 512 routed experts are 708 MB (675 MiB) that only the **decode** reads - a prompt's passes run the draft layer's dense tensors and its K/V - so while a request that **reads** enough tokens is being read, that block can become the prompt path's ring. Before: the ring is whatever the layout could free elsewhere. After: `STRATA_MTP_ROOM_MIN=TOKENS` (opt-in) lends the block in place for such a request and takes it back before the first decode window; a turn that reuses a cached prefix and reads only a few hundred tokens does not take the room or pay its two 675 MiB DMAs. ## Related discussion The original pull request and its discussion: #2830 (closed when its fork was deleted; this PR replaces it). The work comes from the Maya checkout (`project-maya`), local commits `7df369c` (`STRATA_GLM_ROOM_MIN`), `d4d0792` (the tail's slots rounded up to the largest blob) and `5ccf9a4` (the in-place shape, for the MTP config's NextN block). Counting what a request reads, rather than the whole prompt, is this port's change, made after the report below. ## Approach `STRATA_MTP_ROOM_MIN=TOKENS` (`--serve` with `--mtp`): a request that reads at least TOKENS prompt tokens (`0`: every request) uses the MTP drafter's 512 routed experts as the prompt path's ring while it is read. TOKENS counts the tokens this request reads, not the whole prompt: a turn that reuses a cached conversation prefix and reads a few hundred tokens gains nothing from the room. The prompt path's ring is pointed **at the block itself**: `slot i = experts_base + (i - rule) * MAXBLOB` (135 slots on a Q8_0 pack, whose largest blob is 4.98 MiB). Opening and closing a room moves the ring count alone (`Prefill::set_ring_room_active`): the block is never freed, never allocated again and never moves, so no prompt path is laid out again, the decode's captured round/step graphs stay valid (nothing is dropped), and what comes back before the first window is the decode's own bytes (`MtpDrafter::experts_reclaim`). In the measured shape's successor this is two 675 MiB DMAs around the read. The room is opened after the reused prefix is known (`read_from`) and immediately before `lend()`, so a relayout still re-carves with the room; a request that must be re-read from 0 (`STRATA_CKPT_REREAD`) counts all of its tokens. Limits, refused with the reason at start: `--serve` only, `--mtp` needed, `--batch` not supported (its slot drafters share the solo drafter's weights). A layer split is fine (the room goes to the prompt path on the drafter's card), `--pipeline-windows` is fine (the room closes before the request's first window), and `--prefill-head-alias` is fine. Unset (the default) `mtp_room_min` is `-1` and not one line of the room runs; the prompt path is laid out exactly as before. The new code is reachable only behind that check. This is by construction, not by a measurement. `=0` keeps the behaviour that was measured: every request takes the room (the 2026-10-19 control did, on 7-token continuations). ## Validation Tested revision: `41bff1d3` (three commits on `477ac13c`; not rebased, and GitHub reports it merges cleanly with current `main` `e88ee39d`). Environment: Linux, CUDA 13.2, `sm_86` + `sm_120`; the Q8_0 4-way rig is 2x RTX 3060 12 GB and 2x RTX 5060 8 GB. **Passed:** - `ninja -C build strata` on the parent engine and on `41bff1d3` - both exit 0; the two binaries differ by this work. - First shape on our Q8_0 4-way rig (2026-10-16): the room opened, the 675 MiB really went into the ring, the block came back before the first window and the decode ran unchanged. It bought no time at 3,080 or 14,859 tokens and cost 3.4-4.0 s per roomed request, which is why the in-place shape replaced it. - In-place shape and read threshold on our Q8_0 4-way rig (2026-10-19): two engines from the same branch (the parent commit against this one), one start each, `STRATA_MTP_ROOM_MIN=1024`. One 5,042-token chat and three continuations that echoed the answer back. The parent lent the room for the whole 5,096-token prompt on continuations that read 61, 7 and 7 tokens (lend ~102 ms + reclaim ~104 ms); this commit lent nothing. Read time 2,054 -> 1,813 ms, 360 -> 153 ms, 344 -> 152 ms; time to first token 2.110 -> 1.869 s, 0.409 -> 0.203 s, 0.390 -> 0.199 s. A `MIN=0` control on this binary still opened the room on the 7-token continuations (time to first token 0.432 / 0.415 s), so `=0` is unchanged and the difference is the room, not the binary. The cold request still took the room in both arms. - Reported by merbanan on an 8 GB RTX 2060 SUPER (2026-10-18; a port onto a 0.1.45-based branch, not run by us): Qwen3.8-Flash-Next IQ3_XXS, `--serve --mtp` (Q2_0 drafter, 675 MiB), `--resident-budget-gib 40`, `--kv k8v4 --kv-resident 20480`, `--prefill auto`. The auto chunk was 6,656 tokens with a 96-slot ring; the room added 303 slots (399 in all), chunk unchanged. `STRATA_MTP_ROOM_MIN=0` against the same build unset (paired A/B, 3 pairs, prompt tok/s): 4K 448.4 -> 527.2 (+17%), 16K 558.9 -> 658.1 (+17%), 32K 584.0 -> 719.7 (+23%); decode moved -0.27 / +1.42 / -0.17 (noise). The lend cost 107 ms per roomed request. `STRATA_PREFILL_RING` alone showed the ring was the limit (192 slots: +9.8% / +6.2% / +2.3%, decode -3.1 to -3.6 tok/s). First-token logits after a 15,521-token prompt were bitwise equal to the unset run (two runs per arm, all four identical), the 48-token greedy replies matched, and the drafter's acceptance after reclaim was normal (27 of 34). **Failed:** - none observed. **Not run:** - HIP and SYCL builds: no toolchain on this machine. `generate.cpp` carries no per-backend code, so both builds compile the same change. - More than one run per timing on our rig, and `--prefill-head-alias`. - A run against a revision based directly on `e88ee39d`: this head is based on `477ac13c` and was not rebased; GitHub reports it merges cleanly. The measurements above are listed with what they were measured on. No blanket speedup is claimed: the knob is how you measure it, the same prompt with the variable set and unset, `PP` line against `PP` line. `7df369c` records the room as not yet measured, with its standing baseline 89.5 s (1,123 tok/s) on its byte-identical 100,535-token read. ## Limits - The ring gain appears where the ring is the limit (the 8 GB card above) and is nil where it is not (our Q8_0 rig). - The read threshold's gain is separate and does not depend on the ring: it removes the room's own DMA from turns that reuse a conversation, about 0.2 s per short turn on our rig. - `--serve --mtp` only; `--batch` is refused because its slot drafters share the solo drafter's weights. - HIP and SYCL are unverified beyond the fact that the file has no per-backend code. 🤖 Generated with [Claude Code](https://claude.com/claude-code)
Sur le site
Liens install, modèles, releases.