Prompt path: a stager buffer's first job of a generation waits for the previous DMA from it (a narrow race with unpinned blobs)

[sergqwer]

Stager::work waits for a buffer's DMA only from a generation's kRing-th job on (j >= kRing). A generation's first kRing jobs memcpy into their pinned buffers straight away.

So the previous generation's last DMAs from those buffers are only known to be done where something else waits for them:

• staged layers (chunks below stream_all_min): the routing sync (cudaStreamSynchronize(m.cs)) before the layer's start();
• stream_all chunks: the sync at the end of a chunk, which runs only with an MTP drafter (on_chunk) or a layer split (on_stage_chunk).

That leaves a window in two cases:

• a stream_all chunk that ends without that sync (no --mtp);
• a ring entry the routing skipped. Its copy was issued, but nothing waits for it, so even a chunk-end sync of the compute stream does not cover it.

In that window, the next chunk's first jobs can overwrite a pinned buffer that a cudaMemcpyAsync on the copy stream is still reading. It only matters with unpinned blobs, which go through the stager: the RAM-budget tier, or a partially pinned arena.

The fix: the event is now waited for on every job. An event that was never recorded, or has already completed, returns at once. For j < kRing, no job of the current generation can have recorded the event, so the wait is exactly the previous generation's last DMA from that buffer.

We found it in an audit of our fork. We have not seen it corrupt a run: it needs that timing and the unpinned path.

Measured:

• IQ2_XS with --mmap-experts --resident-budget-gib 16: 12,144 of the 48,668 streamed experts went through the stager.
• The tokens on a 32K prompt are the same as on 0.1.31.
• Prefill took 6,546 ms vs 6,491 ms (noise).

🤖 Generated with Claude Code

https://claude.ai/code/session_01VZy1yKaDDiA8a7svdwaHio

[sergqwer]

Rebased on 0.1.32 (c499bd1), no conflicts.

Checked on 0.1.32: IQ2_XS (ISTA), RTX 5090, a fixed --expert-cache, --vram-reserve-mib 1500. The first generated token's logits (STRATA_DUMP_FIRST_LOGITS, now in main) are byte-identical to main's after a 2K and a 32K prompt. Every target builds with STRATA_BUILD_TESTS.

🤖 Generated with Claude Code

https://claude.ai/code/session_01VZy1yKaDDiA8a7svdwaHio

[BlueKingMuch]

Data point for a setup that really streams through the stager: on this 64 GB Windows 11 box cudaHostRegister of the whole 46.8 GiB arena fails and 47 slices (45 GiB) are pinned, so the rest of the experts goes through the stager (RTX 4080 SUPER 32 GB, IQ3_S, 262K, --mtp, --prefill 32768). With this PR merged together with #372, #374, #378 and #407 on 0.1.32: the deterministic decode requests (greedy and sampled, 256 tokens) and the answers after 32K and 64K prompts give the same tokens as before, and the prompt path got faster overall (from #372 and #374), so the extra wait does not show here.

[sergqwer]

@BlueKingMuch thank you for checking this on a setup that really streams through the stager.

🤖 Generated with Claude Code

https://claude.ai/code/session_01VZy1yKaDDiA8a7svdwaHio

[Niko1221]

Released in 0.1.38: the stager waits for a buffer's previous DMA on every job. Merged under your name (GitHub doesn't show it as merged because it went in by hand). Thank you!