#1379 prefill: STRATA_PREFILL_CPU_SHARE=auto times layers with and without the share and shares only while that is faster (#1282)

本文

Follow-up to #1282 (in 0.1.40.2 as an opt-in). **@jase100k measured `STRATA_PREFILL_CPU_SHARE` on an RX 7900 GRE with a Ryzen 7 5700X3D (ROCm), and every share was slower than off:** 0.10 by 2.1%, 0.20 by 3.5%, `auto` (0.47) by 4.5%. Their report also had the number that explains it. This PR makes `auto` check that sharing pays before it keeps sharing.

### Why `auto` lost

`auto` balances two measurements: the CPU's time per expert (c) and the GPU's time per streamed expert (g). It gives the CPU g / (c + g) of the experts, the point where both sides end together. Both c and g are measured with the share on, so nothing compares the layer against the same layer without the share. The CPU's share is not free for the streaming side, and jase100k's numbers show by how much:

| | streamed | host time | host per streamed expert |
|---|---|---|---|
| off | 7,789 | 832.6 ms | 107 us |
| `auto` (0.47) | 5,315 | 1,060.8 ms | 200 us (+87%) |

With the CPU busy, issuing each streamed expert cost the host 87% more. At a share of 0.47 the streaming side keeps 0.53 of the experts at 1.87 times the cost each: 0.53 x 1.87 = 0.99. So the share saved the streaming side nothing. The CPU's work, the copy of the layer's input to the host, the CPU's rows back to the GPU and the join came on top.

I reproduced the rise here by emulating a host-side H2D copy. The emulation is a test patch, below, that memcpys each pinned blob into a pinned bounce slot before the DMA. It ran with the AVX2 kernels and 4 cores / 8 threads (affinity `0xFF`, 7 workers), with a memory-bandwidth hog on the other CCD. Host time per streamed expert rose from 70 us to 130 us (+86%) with `auto` on. The CPU here is fast enough to take 0.73 of the experts, so 0.27 x 1.86 = 0.50 and sharing still won. Whether it pays depends on that product, and only a measurement without the share can tell.

### The gate

Every layer where the share can apply is timed with CUDA events, from before its routing sync to its combine. That window holds everything the share changes: the input copy to the host, the GPU's experts, the wait for the CPU, and the CPU's rows going back. The time is divided by the layer's non-resident experts.

- **Exploration.** The eligible layers alternate between no share and the measured share, starting without, until three ratios are in. A ratio compares two adjacent layers of the two arms: the time per expert with the share divided by the time without.
- **Decision.** `auto` shares while the median of the last five ratios is below 1.
- **Probing.** Every 29th layer then runs the other arm, so both arms stay measured. 29 is prime, so the probes move across the layers from one chunk to the next.
- **Readings it skips.** A layer whose window held one-time costs gives no reading: the host buffers' first allocation or growth (one 250-token run read 6.49 ms where the next read 2.75), and the model's first eligible layer (2x the others' cost per expert in both arms, in every run).

**Why adjacent layers and not a running mean per arm.** I forced each arm on every layer and timed the windows. The window time per expert differs up to 8x between layers (layer 8 takes 30 ms where most take 4-6), and each layer's time repeats within a few percent from run to run. A mean per arm therefore follows whichever layers each arm happened to land on: the first version of this gate did, and it settled on "off" where sharing won by 10%. The same forced runs show the windows hold the whole effect. At 600 tokens the sum of the windows moved by -51.0 ms with the share against -51.4 ms of prompt time (DMA), and by +43.0 ms against +45.6 ms (the one-core case below).

**How the constants were chosen.** I replayed the schedule on those forced-arm windows with 4% noise per reading, 300 runs each. "Every layer" is `auto` before this PR. The last two columns scale the effect down to 0.3, about the size of the report's loss.

| case | every layer | gate | every layer, effect x0.3 | gate, effect x0.3 (worst) |
|---|---|---|---|---|
| DMA 250 / 600 / 900 | 0.814 / 0.821 / 0.885 | 0.822 / 0.821 / 0.883 | 0.944 / 0.946 / 0.966 | 0.948 / 0.959 / 0.970 (0.995) |
| host copy 900 | 0.807 | 0.815 | 0.942 | 0.961 (0.996) |
| one core 250 / 600 / 900 | 1.184 / 1.151 / 1.201 | 1.024 / 1.032 / 1.033 | 1.055 / 1.045 / 1.060 | 1.013 / 1.010 / 1.016 (1.048) |

**What the gate costs.** `auto` cannot know whether sharing pays on a machine until it has tried it. That trial is now bounded:

- **In a one-prompt run** (as in the report), the exploration's 3-4 layers plus one probe. That is a sixth to a quarter of what sharing every layer would cost (the replay's one-core rows).
- **In a running server**, the exploration happens once per process. After that, one layer in 29 runs the losing arm.

### Measured

RTX 5090 + Ryzen 9 9950X3D, IQ2_XS, `--expert-cache 12000`. Prompt ms, 3 interleaved rounds of off / `auto` before / `auto` after; the median is in parentheses. "CPU" is the experts the CPU took over a run.

**DMA** (the default path; every streamed expert is DMA'd from the pinned arena):

| tokens | off | `auto` before | `auto` after |
|---|---|---|---|
| 250 | 397 / 391 / 391 (391) | 356 / 350 / 353 (353, -9.8%), CPU 2,574-2,748 | 364 / 357 / 355 (357, -8.8%), CPU 2,486-2,497 |
| 600 | 505 / 507 / 507 (507) | 457 / 454 / 460 (457, -9.9%), CPU 3,757-3,762 | 461 / 462 / 458 (461, -9.1%), CPU 3,419-3,439 |
| 900 | 577 / 580 / 580 (580) | 523 / 520 / 515 (520, -10.4%), CPU 4,091-4,095 | 520 / 530 / 518 (520, -10.5%), CPU 3,717-3,749 |

**Host copy** (the emulated host-side H2D: bounce memcpy, AVX2, 4C/8T with 7 workers, bandwidth hog). The last column is host time per streamed expert.

| tokens | off | `auto` before | `auto` after | host per streamed expert: off / before / after |
|---|---|---|---|---|
| 250 | 568 / 598 / 569 (569) | 478 / 466 / 462 (466, -18.0%) | 478 / 526 / 469 (478, -16.1%) | 69 / 150 / 124 us |
| 600 | 813 / 813 / 824 (813) | 696 / 678 / 676 (678, -16.7%) | 686 / 690 / 708 (690, -15.2%) | 70 / 130 / 116 us |
| 900 | 921 / 924 / 922 (922) | 789 / 792 / 783 (789, -14.4%) | 789 / 803 / 802 (802, -13.0%) | 70 / 119 / 112 us |

Where sharing wins, the gate keeps about 90% of the gain (89-101%).

**A case where `auto` loses here.** I could not make this machine lose by as much as the report. The closest is the whole process on one core, two threads (affinity `0x3`). That case is noisy: its result moves with whatever else the machine is doing. In this session, `auto` before measured from -4.7% to +11.5% against off. The gate measured from -1.4% to +2.3%:

| one core | tokens | off | `auto` before | `auto` after |
|---|---|---|---|---|
| 1 worker | 250 | 405 / 404 / 400 (404) | 410 / 409 / 420 (410, +1.4%) | 422 / 412 / 403 (412, +2.0%) |
| | 600 | 535 / 529 / 521 (529) | 542 / 536 / 543 (542, +2.4%) | 543 / 539 / 539 (539, +2.0%) |
| | 900 | 600 / 611 / 620 (611) | 613 / 709 / 623 (623, +2.0%) | 625 / 634 / 604 (625, +2.3%) |
| 2 workers | 250 | 414 / 420 / 426 (420) | 425 / 400 / 400 (400, -4.7%) | 415 / 425 / 396 (415, -1.4%) |
| | 600 | 525 / 530 / 533 (530) | 520 / 518 / 551 (520, -2.0%) | 523 / 518 / 554 (523, -1.4%) |
| | 900 | 638 / 613 / 622 (622) | 630 / 631 / 625 (630, +1.2%) | 627 / 621 / 619 (621, -0.2%) |
| 1 worker, earlier in the session: every layer shared, as `auto` before does (one run each) | 250 / 600 / 900 | 413 / 538 / 617 | 447 / 583 / 688 (+8.2 / +8.5 / +11.5%) | |

In the one-worker runs the gate stopped sharing after its exploration in 4 of 9 runs (108-180 experts on the CPU), part of the way in 2 (821 and 944) and kept sharing in 3 (1,086-1,742, where `auto` before took 1,342-1,961). The +2% that remains is the exploration's layers. In these runs `auto` before had settled at a share that cost about as much. I expect the report's machine to look like the "effect x0.3" column of the replay. I cannot run that machine, though, so @jase100k's numbers would settle it.

### Identity

The default (unset) path is unchanged. Against the 0.1.40.2 engine (this branch is on main `e8ca9afd`), at IQ2_XS 2K with `--expert-cache 12000`, the first-token logits are the same bytes and the 32 greedy tokens are the same, 3 of 3 pairs. A fixed share is unchanged too: `STRATA_PREFILL_CPU_SHARE=0.4` at 600 tokens gave the same logits in 7 of 8 runs across both builds. The eighth differed by at most 0.19 with the same tokens, which is the engine's own run-to-run drift: 0.1.40.2 itself gave one such run in 7 at 2K. With `auto`, which layers share now differs, so the output differs from `auto` before in the last bits, as `auto` already did against off.

`STRATA_DBG_CPU_GATE=1` prints each reading and the decision, one line per layer. For example:
`cpu gate: layer 5 shared, 58 non-resident, 2.89 ms (0.0499 a expert) -> share`.

<details><summary>The host-copy emulation (a test patch, not in this PR)</summary>

In `stage_one`, the pinned branch:

```cpp
static const bool bounce_on = [] { const char* v = std::getenv("STRATA_DBG_STREAM_BOUNCE"); return v && v[0] == '1'; }();
static uint8_t* bounce = nullptr;
static size_t bounce_slot = 0;
if (pinned && bounce_on && bounce == nullptr) {
    bounce_slot = (size_t) MAXBLOB();
    if (cudaHostAlloc((void**) &bounce, bounce_slot * STAGE, cudaHostAllocDefault) != cudaSuccess) { err = "..."; return false; }
}
if (pinned && bounce_on) {
    uint8_t* hb = bounce + (size_t) sl * bounce_slot;
    if (m.stage_live[sl]) cudaEventSynchronize(m.copied[sl]);   // its last DMA is done
    std::memcpy(hb, b, (size_t) lay.blob_bytes(l));
    if (m.stage_live[sl]) cudaStreamWaitEvent(m.copy, m.used[m.used_of[sl]], 0);
    cudaMemcpyAsync(m.stage_dev[sl], hb, (size_t) lay.blob_bytes(l), cudaMemcpyHostToDevice, m.copy);
    ++stats_.experts_dma;
} else if (pinned) { /* as before */ }
```

The runs used `STRATA_DBG_STREAM_BOUNCE=1 STRATA_FORCE_AVX2=1 --pool-workers 7` under affinity `0xFF`, with 6 memcpy threads on CPUs 16-31 while each run held the GPU.

</details>

🤖 Generated with [Claude Code](https://claude.com/claude-code)

https://claude.ai/code/session_01VZy1yKaDDiA8a7svdwaHio