Pull requests / #3483
#3483 [Engine] stage buffers: fence a queued copy per stream, not per buffer (#1237)
open · @John-Pi · 0 comments · View on GitHub
Server & APIAMD / HIPNVIDIA / CUDAModels & quantsDocumentationWindowsLinux
Description
## Problem and result
With `STRATA_STAGE_PIN=1` the stage buffers are page-locked, so the cache fill's `cudaMemcpyAsync` is really asynchronous: the DMA reads the buffer after the call returned. `f71c5307` fenced that by recording an event after each queued copy and having `claim_stage` wait for it before recycling the buffer — but **one event per buffer**, and a buffer can have copies queued on more than one stream (the hit path, a peer GPU's refill, the startup's refills all read the same blob). The second `cudaEventRecord` overwrote the first, so at recycle time `cudaEventQuery` answered for the stream that was already done while the slower stream's copy still had the buffer under the DMA. The refill then overwrote it — the corrupted IQ3_S decode of #1237 (48 tokens of id 0 after a 4096-token prompt).
Before: a refill can run under a queued copy on another stream. After: the refill waits for **every** copy queued out of the buffer, whatever stream it left on.
## Related discussion
#1237 is the pinning PR this came out of, and it is where the recycle path was named as the blocker:
> It stays opt-in because pinned buffers made the cache-fill copy asynchronous and a stage buffer could be reused before its copy had finished; **a fix of the recycle path in `src/core/expert_source.cpp` would let us make it the default.** — @Niko1221, 2026-10-17
This is that recycle path. It is a bug fix in code that is already in `main`, so no feature scope to agree on; **`STRATA_STAGE_PIN` stays opt-in and no default changes** (see Limits).
## Approach
- **One fence entry per stream.** `FileExpertSource::StageFence { stream, event, live }`, a fixed table of 4 per buffer. `note_async_read` reuses the entry of the same stream if there is one (the stream's own order makes the later copy the last to finish), otherwise takes a free entry; if the table is full it waits for the oldest entry and reuses it, so a fence is never dropped.
- **`claim_stage` waits for every live entry** before it hands the buffer out again, and counts those waits in `stage_fence_waits_` (`stage_fence_waits()` is new; the test reads it).
- **`close()` and `end_startup_unbuffered()`** synchronise before destroying events, as before; the startup path now also drops the fence table with the buffers.
- **Pageable buffers are untouched**: the driver stages them before the call returns, so they still have no event.
- The victim choice is unchanged (least recently used), with a comment saying why: a two-pass "prefer a buffer no copy is reading" was measured on a Radeon AI PRO R9700 and cost 59% more file reads and 20% of decode.
- `tests/hip/stage_fence.cpp` is the deterministic reproduction, registered as ctest `hip_stage_fence` (skips with 77 without a device or without unbuffered reads).
The test builds a synthetic 4-layer × 2-expert pack in the working directory, takes the (0,0) blob into a stage buffer, queues one copy of it behind ~4 GiB of device copies on stream A, queues the same blob on stream B (which completes at once, and whose record is the one a single shared event keeps), lets `kStageAge` layers pass so the buffer is recyclable, then refills the same buffer. It asserts the refill waited, and that both copies still hold the (0,0) bytes.
## Validation
**Passed:**
- `cmake --build build-hip -j$(nproc)` — 338 of 338, Release, `CMAKE_HIP_ARCHITECTURES=gfx1201`, RX 9070 (gfx1201, 16 GB), system ROCm 7.2.4, kernel 7.2.11, engine 0.1.46.1 at `9ae93f7b` + this change.
- `ctest --test-dir build-hip -j4` — 78 of 80. `hip_stage_fence` passes in 0.16 s and 3 of 3 standalone runs:
```
FileExpertSource: the stage buffers are pinned (cudaHostAlloc)
stage_fence: two queued copies, one refill, 1 fence wait(s), 18.0 ms; both copies hold the pre-refill bytes
```
**Failed:**
- `hip_stage_fence` **without** this change (same tree, `9ae93f7b`, only the test and the `stage_fence_waits()` accessor added) fails as designed, which is the reproduction:
```
stage_fence: the stream A copy landed the REFILLED bytes (marker 8, expected 1): a
page-locked stage buffer was recycled under a queued copy (#1237)
```
exit 1. That is the bug: stream B's already-complete copy hid stream A's queued one.
- `ctest`: `ple_parity` (wants `../../Q2_0/Qwen3.8-Flash-Next-GSQ-RCO-Q2_0-00002-of-00002.gguf`, not on this machine) and `hip_prefill_mmq_parity` (its APR block calls `ggml_type_name(145)` while the TQ overlay leaves `GGML_TYPE_COUNT` at 43, so the read is out of `type_traits[]`; `basic_string: construction from null is not valid`). Both fail identically on `origin/main` in this environment, and neither file is touched here.
**Not run:**
- **CUDA and SYCL builds.** `src/core/expert_source.cpp` is shared with the CUDA build and I have no `nvcc` on this machine; SYCL is untouched by the change. The diff is host-side `cudaEvent*`/`std::array` only, but it is unbuilt there.
- **An end-to-end engine A/B of the fence.** Showing the corruption return would need a 4096-token IQ3_S prompt on a pageable-vs-pinned arm; the test reproduces the same race deterministically, which is what I ran. I did run the engine with `STRATA_STAGE_PIN=1` (1 199-token prompt, `--mmap-experts` resident mode): normal prefill and decode, `stage_fence_waits()` 0 for the run.
- Windows.
## Limits
- **The default is not flipped.** `STAGE_PIN_DEFAULT` stays 0 and `docs/DETAILS.md` keeps saying opt-in. Flipping it is a separate decision with its own numbers; the pinned-vs-pageable decode measurements (R9700, 6 pairs, +2.1% to +3.5%) belong to the pinning change, not to this one. Say the word and I will open that as its own PR.
- **4 entries per buffer** is a bound, not a measurement: more than 4 streams copying out of one stage buffer at once waits early rather than dropping a fence. In the engine at most 3 streams appear (hit path, peer refill, startup).
- **The test needs a real device and an unbuffered file tier** (`O_DIRECT`), so it skips on tmpfs and on machines without a GPU; it is not a substitute for the CUDA-side coverage.
Related on this site
Quick links to related pages.