Pull requests / #269
#269 Split prefill: run the whole prompt on the main GPU (--prefill-main)
closed · draft · @mijkathegreat · 0 Kommentare · Auf GitHub
BenchmarksServer & APIMulti-GPUNVIDIA / CUDAModels & quantsWindows
Beschreibung
## 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.
Mehr auf der Site
Links zu Install, Modellen, Releases.