#2191 --prefill-head-alias (opt-in, --serve, one GPU): the output head's VRAM holds the prompt buffers while a prompt is read

Descripción

NOTE: This is human-designed, AI-coded, we optimised memory allocation per phase, and achieved substantial measured boost in decode and prefill speed. Please check the draft, hope it helps at least as an inspiration, I might not have time to polish it up, will try to answer questions. Thanks for considering! Max.

## Title

`--prefill-head-alias` (opt-in, `--serve`, one GPU): the output head's VRAM holds the prompt buffers while a prompt is read

Issue: none. I searched open issues and PRs for head/prefill-workspace/phase-VRAM work and found nothing related.

## Summary

On an 8 GB card the prompt chunk is as large as the free VRAM allows, and prefill speed follows the chunk closely: 185 tok/s at 256 tokens, 292 at 512, 443 at 1,024 and 631 at 1,792, all on the same 29,854-token prompt. Without this change, every chunk above 768 tokens needed a lower `--vram-reserve-mib`, and that left 0 MiB free, which is where WDDM starts paging.

The idea is to treat VRAM as phase-dependent storage. Two large users never work at the same time:

- the **native output head** (417 MiB, Q5_K on IQ3_XXS) is needed only while generating;
- the **batched prompt path's buffers** are needed only while a prompt is read.

With `--prefill-head-alias`, the head is unloaded while a prompt is read, its VRAM is used for prompt buffers, and it is read back in before the first decode window. The auto cache is also sized with the head out, so the prompt path can borrow almost the whole cache (slots 1..N) for a 2,048-token chunk. After the prompt, the cache is refilled as usual.

**Results.** All numbers are from an RTX PRO 2000 Blackwell Laptop (8 GB, Windows 11, WDDM), IQ3_XXS, 590 MiB reserve, on a local 0.1.41 build with this change. The full tables and settings are in `docs/PREFILL_HEAD_ALIAS.md`.

| | Default (auto chunk 512) | `--prefill-head-alias` (chunk 2,048) |
|---|---:|---:|
| Prefill, 29,854-token prompt, fresh engine per trial | 292.5 tok/s (12 trials) | 667.1 tok/s (6 trials) |
| Expert-cache slots | 508 | 749 |
| Free VRAM after start / while decoding | 287 / 207-287 MiB | 289 / **0 MiB** |
| Agent loop, prefill (context-matched) | 256 tok/s | 583 tok/s (+125%, 95% CI +123 to +129) |
| Agent loop, decode (context-matched) | 24.8 tok/s | 29.2 tok/s (+19%, 95% CI +17 to +23) |
| Agent loop, expert-cache hits | 20.8% | 25.9% |

The agent-loop rows come from 1,200 requests over 8 cycles (default) and 2,649 requests over 20 cycles (head alias) of OpenCode agents. A large part of the decode difference comes from the default runs' long slow stretches near 19.5 tok/s, which the head-alias run did not have. Against the default's fast stretches only, decode was about 6% faster. We do not know what caused the slow stretches.

A controlled replay held the prompt and a 3,072-token forced output fixed, with a fresh engine per trial:

| Setting | Prefill | Decode |
|---|---:|---:|
| 1-slot cache | 441 tok/s | 22.3 tok/s |
| 1-slot + head alias | 666 tok/s | 24.5 tok/s |
| auto cache + head alias | 668 tok/s | 26.2 tok/s |

The default (auto cache without the alias) was not run in that forced mode.

**Cost.** Short reads get slower. Every request reloads the head (median 255 ms) and lends and refills almost the whole cache. The engine's prompt time (production medians, default -> head alias) by fresh tokens read:

| Fresh tokens | Default | Head alias |
|---:|---:|---:|
| 1-64 | 671 ms | 1,311 ms |
| 65-512 | 1,703 ms | 2,453 ms |
| 513-2,048 | 4,062 ms | 3,331 ms |
| 2,049-8,192 | 15,206 ms | 6,876 ms |

It pays off from somewhere between 513 and 2,048 fresh tokens.

**How this was made.** The direction and the idea are mine: use VRAM as phase-dependent storage, and chase bigger prompt chunks because our benchmarks showed a clear trend. I understand Strata's memory layout only in part. The code was written with an AI coding assistant ("vibe-coded"), then measured and run in production on my machine for about 26 hours (2,717 requests) with no errors or restarts. Please review it with that in mind. The areas I am least sure about are listed under Extra Notes.

## What changed

- `src/program/generate.cpp`
  - `--prefill-head-alias` option, usage text, and refusal reasons: needs `--serve`, one GPU and no layer split, no `--batch`, no `--pipeline-windows`, `--short-read 0`, `--expert-cache 1|auto`, `--mtp`, and a native head.
  - At start: MTP binds its draft head from the output head; then the head is unloaded before the auto cache is sized. The loan starts at slot 1, and the chunk is the largest of 2,048/1,536/1,024/512 that fits the borrowed slots plus free VRAM with 512 MiB to spare. A `cudaGetLastError()` retry covers a 0-byte free reading right after the unload.
  - Per request: unload the head and rebuild the prompt path before reading; free the prompt buffers and reload the head after refilling the lent slots.
  - The per-request loan keeps the start-up layout in alias mode. The upstream #2090 sizing is unchanged when the flag is off; it only moved inside `if (!o.prefill_head_alias)`.
  - The verifier is given the (temporarily unloaded) head, and MTP is not bound a second time.
  - Guard: the verify graphs captured in the first decode hold the head's address. If a later reload lands elsewhere, the engine stops with `ERR` (and the server restarts it) instead of running stale graphs.
- `include/strata/core/native_head.hpp`, `src/core/native_head.cpp`: `NativeHead::unload()`, which the destructor now uses.
- `include/strata/core/mtp.hpp`: `MtpDrafter::set_window_R()`, so the verify rows can be attached after an early bind.
- `include/strata/prefill/prefill.hpp`, `src/prefill/prefill.cpp`: `Prefill::init(..., bool borrow_spill = false)`. When the borrowed region is full, `Alloc::take` continues with owned `cudaMalloc` buffers (the "hybrid" layout). `relayout` refuses a hybrid layout. With `borrow_spill` false, `take` behaves exactly as before.
- `sycl/src/prefill/prefill.cpp`: the matching `init` signature. It refuses `borrow_spill`, because the SYCL engine has no flag for it.
- `docs/PREFILL_HEAD_ALIAS.md` (new) and a bullet in `docs/DETAILS.md`.

What it leaves alone: the default path. Every new branch is behind the flag or behind `borrow_spill`. No defaults changed, and there are no server, setup or API changes. To use it, add `--prefill-head-alias` and `--short-read 0` to the model config's `args`.

## Extra Notes

**How I checked:**
- Built this branch (rebased on `main` @ a0116b27, v0.1.43.1) for CUDA 13.0, sm_120, MSVC 14.44 on Windows. It compiles.
- Smoke-tested this branch on the same PC, flag off and on, with the same build and a fresh engine each. Both runs used `STRATA_ARENA_PIN_GIB=auto`; without it, the full sliced arena pin left 0 GiB free for the cache on this 64 GB PC (#243), with the flag off as well. The requests were a 29,842-token prompt cold, a 31-token follow-up on its prefix, two 45-token prompts, and the long prompt again. All completed with coherent answers. The cold prompt read at 276.5/277.1 tok/s with the flag off and 643.1/649.7 tok/s with it on; decode was 35.7/34.7 and 34.1/31.6 tok/s (free generation, 768 tokens). The head came back at the same address on every reload.
- The longer measurements (tables above) come from my 0.1.41 build with the same change. The upstream #2090 loan sizing was the one real merge conflict.
- Default path: by construction (see What changed). I have not run window-hash parity of flag-off against `main`.

**Not tested:** HIP (same sources; I have no AMD card), SYCL (signature only), Linux, other cards and models, multi-GPU (refused).

**Open questions for review:**
1. **Head address.** `cudaMalloc` does not promise the same address on reload. The guard turns a moved head into a restart, but a cleaner fix would keep the address: a CUDA VMM range whose backing is swapped, or prompt buffers placed inside the head's own allocation instead of freeing it. Which do you prefer?
2. **0 MiB free while decoding.** The cache is sized with the head out, so after the reload the verify windows were captured with 0 MiB free on this card (the default had 207-287 MiB), and 0.1.43.1's new start-of-request warning fires at 0 MiB. We saw no paging over 2,717 requests, but it may be wrong for other cards. Should the sizing reserve the head's bytes for decode and spend only the prompt path's own share?
3. **Short reads.** A threshold on fresh tokens could keep the head and read through the decode windows, as `--short-read` does, removing the cost above for small reads: 614 of 2,717 requests read 64 or fewer.