Pull requests / #269

#269 Split prefill: run the whole prompt on the main GPU (--prefill-main)

closed · draft · @mijkathegreat · 0 comentarios · En GitHub

BenchmarksServer & APIMulti-GPUNVIDIA / CUDAModels & quantsWindows

Descripción

## One GPU outruns the split on prefill

Same rig, same model (IQ3_XXS), same calibrated config, greedy, fresh process per run:

| prompt path | tok/s prefill | tok/s decode |
|---|---:|---:|
| single GPU (all 48 layers on CUDA0) | **1704** | 59 |
| layer split, stock (auto K=24) | 990 | 66 |

Test rig: 2× RTX 5090 Laptop — one internal on **PCIe 5.0 x8** (26.7 GB/s measured to the pinned host arena), one in a **TB5 enclosure** (Gen4 x4, 5.3 GB/s); 64 GB RAM, Windows 11 / WDDM. Prompts of 15.5k–20k tokens.

## Why the split loses to one card

The batched prompt path is bound by **streaming expert misses from the host arena, not by compute** — one GPU already finishes an 8192-token chunk faster than its misses can be delivered. A layer split does not split that stream; it adds obstacles:

1. **The second card streams its own layers' misses over its own link.** Half the layers depend on the TB5 card's 5.3 GB/s, which puts the slower link into the critical path of every chunk.
2. **Windows caps arena registration at 8 GiB once a second context exists** (single-GPU mode registers the whole ~39 GiB arena). Most misses then move through host-copy staging instead of DMA.
3. **The stages serialize**: chunk hand-offs go through pinned host memory (no P2P through WDDM), and each stage carves its own chunk-sized prompt buffers out of what would otherwise be expert-cache VRAM.

So the second GPU does not add bandwidth where prefill needs it — it adds a slow link and takes away the pin.

## This PR: prefill everything on the main GPU

Same binary, same config — `STRATA_PREFILL_MAIN=1` (+ `STRATA_PIN_LIMIT_GIB=20`) vs off:

| layer split | tok/s prefill | tok/s decode |
|---|---:|---:|
| stock (this PR, env off) | 990 | 66 |
| **this PR, env on** | **1520** | 75 |

Up to +54% prefill over the stock split, decode unchanged or slightly better (the second card's VRAM goes to its expert cache instead of prompt-path buffers). Output verified clean: byte-identical greedy continuations on short prompts, correct long-prompt answers, decode cache hit 98.9%, no non-finite values under `STRATA_DBG_NAN`.

The remaining gap to single-GPU (1520 vs 1704) is *not* the prompt path itself — that is now literally the single-GPU path — but the memory system around having a second context: the Windows registration budget tops out at ~20–22 GiB with two contexts (single-GPU mode pins the whole ~39 GiB arena, and past the ceiling registrations fail and poison the contexts); the main cache's profile covers only its own layers `[0, K)`, so the later layers stream cold during the prompt; plus the hand-off copy per segment.

## Principles of the solution

1. **All layers' batched prefill runs on the main device** (`sp.set_stage(0, -1, nullptr)`); the stages skip their prompt-path init and keep that VRAM for their caches. Decode stays stock-split — only the prompt path moves.
2. **One hand-off per prompt segment**: `session_copy_layers()` copies layers `[K, n_layers)` of the session — GDN rows, whole QSA carve blocks, and the streamed KV host pools — with `cudaMemcpyPeer` hops. A plain two-step pinned staging copy *silently lands nothing* across Thunderbolt on WDDM (destination reads back as zeros after a "successful" copy), so the driver-owned peer path is used directly.
3. **Residency must be masked in place, never snapshotted**: the later stages' slot ids index *their own* device's arena, and the shared expert pool reads the same `host_res` array. During the batched prompt those ids are masked out in `host_res` itself, and a per-request snapshot restores the true table before the verify windows. A stale copy of the table once made the prompt path read its *lent workspace* as expert weights — NaN from the first QSA layer onward; that is the bug class this guards against.
4. **`STRATA_PIN_LIMIT_GIB`** raises the Windows multi-GPU registration cap. 20 GiB is measured safe on 64 GB RAM; 22 is the ceiling — above it, registrations fail and poison the contexts.
4. **Which card is the main one matters** (the path is stream-bound): with the Gen5 x8 card as CUDA0 the split prefills at 1384 tok/s, with the TB5 card as CUDA0 at 890. nvidia-smi numbering floats between sessions, so the second commit makes the server order the split's GPUs by PCIe link class and put the widest one first (logged when it reorders; falls back to the configured order if the probe fails).

Both changes are off by default: without the environment variable the binary is byte-for-byte stock behavior.

En el sitio

Enlaces a install, modelos, releases.