Pull requests / #3105

#3105 Add request-scoped prefill VRAM loans

open · @c3c · 0 comentarios · En GitHub

BenchmarksServer & APIAMD / HIPNVIDIA / CUDAModels & quantsWindows

Descripción

This the rebased version of https://github.com/Niko1221/Strata/pull/1581 for the VRAM loaning part of the PR.

## Summary

This PR adds a manually enabled, request-scoped VRAM loan for prompt prefill.
It temporarily lends decode-only allocations to the prompt path, then restores
the verifier, native output head, MTP experts and draft head before generation.
The repeat-prefill option applies the same transaction to later long messages.

The default path is unchanged when the flags are absent.

## What changed

Enable the feature with:

```json
{
  "STRATA_PREFILL_ELASTIC_LOAN": "1",
  "STRATA_PREFILL_MTP_LOAN": "1",
  "STRATA_PREFILL_HEAD_LOAN": "1",
  "STRATA_PREFILL_RETAIN_STARTUP_CHUNK": "1",
  "STRATA_PREFILL_REPEAT_LOAN": "1",
  "STRATA_PREFILL_REPEAT_MIN_TOKENS": "50000"
}
```

All six `STRATA_PREFILL_*` variables above are introduced by this PR, but they
are not six independent requirements: `STRATA_PREFILL_ELASTIC_LOAN` is the
master switch; MTP/head loans are staged sub-options; and the retain/repeat
flags apply only when later prompts should repeat the transaction. A minimal
startup experiment needs only the master switch. The existing upstream
`--vram-elastic` command-line option is required as the segmented-cache
backend; it is not part of this PR.

| Variable/option | Status |
| --- | --- |
| `STRATA_PREFILL_ELASTIC_LOAN` | New master switch |
| `STRATA_PREFILL_MTP_LOAN` | New optional MTP/draft deferral |
| `STRATA_PREFILL_HEAD_LOAN` | New optional native-head deferral |
| `STRATA_PREFILL_RETAIN_STARTUP_CHUNK` | New optional retention of startup geometry |
| `STRATA_PREFILL_REPEAT_LOAN` | New optional subsequent-prompt transaction |
| `STRATA_PREFILL_REPEAT_MIN_TOKENS` | New repeat threshold; default 1,024, minimum 256 |
| `--vram-elastic`, `--serve`, `--mtp`, `--spec` | Existing upstream options used/validated by this PR |

`STRATA_PREFILL_REPEAT_MIN_TOKENS` applies only to subsequent prompts. It is
clamped to a minimum of 256 tokens and defaults to 1,024. Startup loans are
controlled by the startup flags because the first request length is not known
when the engine initializes.

The implementation is in the serve/prefill lifecycle (`src/program/generate.cpp`)
and the verifier, native-head and MTP resource transactions
(`src/core/verify.cpp`, `src/core/native_head.cpp` and `src/core/mtp.cpp`).

As mentioned in the original PR: **benefit depends on prompt length and whether the released VRAM
crosses a useful planner threshold**. 

## Extra Notes

### Test system and method

- GPU: NVIDIA GeForce RTX 2080 Super Max-Q, 8 GiB, 105 W, compute capability
  7.5.
- CPU/RAM: Intel i7-10875H, 64 GiB RAM.
- Link: PCIe 3.0 x16; measured host-to-device bandwidth approximately
  12.2 GB/s.
- Platform/build: Windows, CUDA 13.4, MSVC Release build with SM75 SASS.
- Engine: current upstream main at v0.1.45 (`477ac13c`) plus this PR.
- Model: Qwen3.8-Flash-Next IQ3_XXS, native GGUF plus IQ3_XXS expert pack.
- Runtime: 65,536-token context, Q4_0 KV, 24,576 resident KV cells, MTP
  enabled with `--spec 4`, and the experimental speed projection held constant.
- Benchmark request: 31,258 fresh prompt tokens and a 512-token completion
  cap. Runs were serialized with a fresh engine; PP/TG are the engine-reported
  timing-line rates.

### Guarding and supported configurations

The feature is restricted at startup to the serialized, one-GPU NVIDIA CUDA
server path. It requires `--serve`, `--vram-elastic`, CUDA virtual-memory
segments of at least 64 MiB, and rejects batch slots, pipeline windows, layer
split, helper/remote/peer caches, HIP and `--prefill-head-alias`. MTP/head
loans require `--mtp` with `--spec >= 2`; the head loan additionally requires
a native output head.
Invalid combinations fail before any decode resource is deferred.

CPU worker count does not change the ownership model. Multiple conversations
can use one server sequentially, but concurrent batch sessions are rejected.
Separate Strata processes sharing one GPU are not coordinated and must not run
concurrently with loans enabled.

### Measurements

The fixed-planner A/B isolates the transaction when geometry is held constant:

| Configuration | Prompt plan | PP | TG |
| --- | ---: | ---: | ---: |
| Loan flags off | 2,304 / 227 | 424.7 | 29.8 |
| Loan flags on, same plan | 2,304 / 227 | 425.1 | 29.3 |

This is effectively neutral (+0.1% PP, within run noise). The feature is not a
same-kernel speedup; its purpose is to make a larger prompt geometry fit.

With automatic planning (`--prefill auto`), the loan-only run selected a larger geometry:

| Configuration | Prompt plan | PP | TG |
| --- | ---: | ---: | ---: |
| Loan flags on, automatic planner | 5,376 / 227 | **780.6** | **29.7** |

The combined configuration is **11.5% faster in PP** than loans alone in this
screen. The planner also changes the chunk/ring geometry, so this is reported
as an end-to-end combined benefit rather than a loan-only percentage.

A repeat-prefill test released and reloaded the verifier, native head, MTP
experts and draft head without a CUDA/runtime error. The loan-only follow-up
completed at **742.9 PP / 30.4 TG**.

This PR is intentionally incompatible with `--prefill-head-alias`; startup
fails before a loan is taken. Use head-alias mode by itself, or remove it when
using the loan transaction.

### Cross-GPU VRAM-loan qualification

The loan is not SM75-specific. A separate test used an RTX 4070 Ti
(12 GiB, SM89), INT8 KV, a 196,608-token context, a 512-slot staging ring - but an older v0.1.41 source revision. Loan-off and loan-on runs were interleaved,
serialized, and started from a fresh engine. The 36k result is a three-run
median and the 72k result is a five-run median; both used a 512-token
completion cap.

| Fresh prompt | Loan | Prompt chunk | PP (tok/s) | TG (tok/s) | GPU prefill |
| ---: | --- | ---: | ---: | ---: | ---: |
| 36,061 tokens | Off | 8,448--8,960 | **2,562.9** | **54.8** | 13.811 s |
| 36,061 tokens | On | 15,104 | 2,539.4 | 52.9 | **12.987 s** |
| 72,061 tokens | Off | 8,960 | 2,617.0 | 54.1 | 27.271 s |
| 72,061 tokens | On | 15,104 | **2,680.9** | **54.9** | **25.772 s** |

The larger loan-backed chunk reduced median GPU prompt work by 6.0% at 36k
tokens and 5.5% at 72k. Loading the deferred output-head and MTP resources is
a fixed boundary cost, so the 36k request was 0.9% slower by the engine's
end-to-end PP rate while the 72k request was 2.4% faster. The five-run 72k PP
ranges were 2,604.3--2,633.7 off and 2,501.1--2,696.3 on; four of five loan-on
runs were 2,673.5 PP or higher, with one retained slow outlier. This is evidence
that the general CUDA loan works beyond Turing, but also that it should remain a
manual opt-in: benefit depends on prompt length and whether the released VRAM
crosses a useful planner threshold.

En el sitio

Enlaces a install, modelos, releases.