[VanshSharmaPES]

Draft, branched from #1178's head (77dc04e). #1178 has to land first, and this branch will need a rebase once it does. Review only the last three commits (2afd7d4, ca8c6e5, d9fc2a1); the rest is #1178. I have not measured whether this is faster, and I can't: it needs the 5070.

This is part 2 of #1174: the zero-VRAM alternative to a second slot (option (c)), to be measured before it.

What I changed

With #1178 the embedding is reloaded at the tail of the previous backward, so what is left of the large-layer wait is the head's forward-tail load into the shared slot. tail_prefetch fires from StreamPrefetcher.advance() at the last decoder layer, before that layer's functional_call, so the copy hides behind one layer of compute.

The embedding's bytes in the slot are dead once the lookup has run, and the lookup precedes layer 0's advance, so I added StreamPrefetcher.head_prefetch_layer and default it to 0, the earliest point. None restores the last-layer timing. It is a plain attribute, so one process can run both timings, the way backward_tail_prefetch=None gave arm A on #1178:

runtime.prefetcher.head_prefetch_layer = None   # arm A: today's timing
runtime.prefetcher.head_prefetch_layer = 0      # arm B: this PR (any layer in [0, n_layers) is accepted)

The refill is ordered after the queued embedding lookup by the wait_stream in LargeLayerBufferPool.load_async, the same guard #1178 relies on.

Evidence (CPU, and an RTX 3050 with torch 2.14.0+cu130)

• Bit-equality. Per-step gradients and the final adapter weights are torch.equal between the two timings for untied SFT (also at gradient_accumulation_steps: 2), ORPO, SimPO, DPO and KTO, and for tied SFT, where there is no large pool and the switch does nothing. The same timing run twice is a control, so the comparison is not noise. I randomise lora_B first, since KTO's gradients in bf16 were all zero without it and the comparison passed while checking nothing.

• Load counts. The real large-layer load sequence and count are identical in both timings: 5 over two steps. The head load's prev is 0 in arm B and the last layer in arm A, so the move really happens. An eval forward followed by a training step gives 2 loads after the eval and 5 after the step, in both timings.

• Mutations of the main change (CPU, 21-test version of the file):

  mutation	result
  none	20 passed, 1 skipped
  head prefetch call deleted from advance	15 failed
  head prefetched in _prime, before the lookup	11 failed
  once-per-pass flag reset removed from prime()	12 failed
  override ignored (always last layer)	4 failed
  range validation removed	3 failed
  wiring reverted in install_streaming	passed at first; I added tests for the installed default, and it now fails 1
  direction check dropped from the trigger	passed: equivalent, the once-per-pass flag already blocks a refire
  idx == head_layer to idx >= head_layer	passed: equivalent, same reason

• Pool-level GPU test where the queued read is the embedding lookup (@pytest.mark.gpu, skipped without CUDA). On the 3050 it passed 5 of 5 with the wait_stream and failed 5 of 5 with it patched out.

• The whole new file passes on CUDA (23 passed), and with tests/test_issue975_backward_tail_embed_prefetch.py it is 28 passed. On CPU: ruff is clean, and pytest -k "stream or 1049 or 975 or 1174 or docs or changelog" gives 624 passed, 36 skipped, 0 failed. mypy reports the same errors as on fix(layer-stream): prefetch embed_tokens at the backward tail, not the next step's prime (#975) #1178's file.

What I checked about your two hazards (not measured)

• Copy stream. On the RAM tier the head's get moves from after layer 3 to right after layer 1's prefetch (4-layer untied fixture), so layer 2's copy now queues behind it, as you predicted. Whether that costs more than it saves needs the GPU.
• Disk tier. I recorded every source.get() over 3 steps on AsyncDiskSource, on a tiny CPU fixture, three runs per arm. The move added exactly one demand miss, on the head at step 0 (arm A: 6 misses, arm B: 7-8). The decoder-layer misses were the same in both arms, and steady state was identical. This says nothing about real timing, and it is why the disk tier has to be measured. The change needs no edit to _plan_queue, _claim_slot or get.

What I have not done

• No measurement. Per your comment on predicted_peak_bytes under-predicts VRAM for an untied streamed head once #1061 lands; is #975's one-slot ping-pong worth trading for 2 slots? #1174, it should run on fix(layer-stream): prefetch embed_tokens at the backward tail, not the next step's prime (#975) #1178's smaller fixture, judged by the step's total wait (every entry in stalls) on both the RAM and the disk tier. If layer 2's wait grows, two other placements are worth trying, and head_prefetch_layer accepts them: a later layer, or a stream of its own for the head.
• No prime() layer-0 change. I left that for its own issue.
• Nothing else is missing for the CPU half; the changelog fragment is in (d9fc2a1).

[MakazhanAlpamys]

@VanshSharmaPES, thank you for turning the go-ahead on #1174 into a prototype within hours, and for bringing all the evidence that list asked for. Three things stand out:

• The switch is a plain attribute that advance() reads on every call, so both timings really can run in one process. I flipped it every step (A, B, A, B) and the head's load landed at layer 3, 0, 3, 0.
• _run_arm randomises lora_B and has a same-timing control (test_the_comparison_is_not_noise), so the bit-equality class cannot pass vacuously. Your KTO note, where the bf16 gradients were all zero, is exactly the reason that matters.
• Your mutation table shows the two survivors and explains why they are equivalent, instead of leaving them out.

What I ran

Everything ran on CPU (CUDA_VISIBLE_DEVICES=-1), with soup_cli imported from a checkout of d9fc2a17. There was no GPU where I checked this.

• Your two test files: 26 passed, 2 skipped (the two GPU tests). ruff check is clean.

• The new file on unmodified main (afd9a917): 14 failed, 8 passed, 1 skipped. On fix(layer-stream): prefetch embed_tokens at the backward tail, not the next step's prime (#975) #1178's head (77dc04e3) it gives 11 failed, 11 passed, 1 skipped.

• Bit-equality with fix(layer-stream): prefetch embed_tokens at the backward tail, not the next step's prime (#975) #1178, across trees. I compared the two trees, not just the two arms inside yours. The same 2-step run through the real setup() path, on 77dc04e3 and on d9fc2a17 at its default timing (head at layer 0), for untied SFT, SFT at gradient_accumulation_steps: 2, ORPO, SimPO, DPO, KTO and tied SFT:

  • 32 of 32 gradients match with torch.equal, all non-zero;
  • 16 of 16 final adapter tensors match;
  • the large-layer load counts are the same (5, 9, 5, 5, 9, 17, and none for tied).

• Mutations of the main change, run against both files:

  mutation	result
  installed default 0 → None (layer_stream_runtime.py:2696)	1 failed
  configured layer ignored, always the last (:1044)	4 failed
  self.tail_prefetch() deleted from advance() (:1057)	18 failed
  prime() no longer resets tail_prefetched (:1028)	14 failed
  head loaded in _prime, before the lookup	14 failed
  constructor range check removed (:1005)	3 failed
  off by one, idx == head_layer + 1	14 failed
  direction check dropped (:1048)	passes, equivalent (as you said)
  == → >= (:1049)	passes, equivalent (as you said)
  wait_stream removed from LargeLayerBufferPool.load_async (:951)	passes on CPU. Only the two GPU tests pin it, and they skip without CUDA.

  Your counts reproduce. Where mine are higher, the difference is the The single shared large-layer slot forces a full embedding-matrix H2D every step on an untied checkpoint (~33 ms for 500 MiB, measured at 16 GB/s) #975 file's tests, which I ran alongside.

Your two hazards, checked on CPU

1. The shared copy stream. I recorded the order in which copies reach the two pools on a 4-layer untied fixture (RAM tier). Both pools share one copy stream on CUDA, so this is the order they would queue in there.

step	arm A	arm B
SFT, steady state	L0, L2, L3, head	L0, head, L2, L3
policy forward after a no-grad forward	—	embed, L0, L1, head, L2, L3

Your reading holds: layer 2's copy queues behind the head's. One addition: at every SFT step start, the head now also queues behind prime()'s copy of layer 0 (:1029), which @MakazhanAlpamys pointed out on #1174 is made even when layer 0 is already resident. Read the layer-0 entry and the head's entry together in the measurement.

2. The disk tier. At the default depth I also saw no difference in decoder misses. At stream_read_ahead: 3, though, the drop described in the go-ahead does happen. The setup:

• a 6-layer untied fixture on AsyncDiskSource with pageable staging;
• every shard read slowed to 20 ms and every decoder layer's compute to 40 ms, so the lookahead is still pending when the head's get arrives;
• every source.get() logged, with the reader's queue before and after.

What it showed:

• Preference-shaped step (a no-grad forward, then forward and backward), read_ahead=3. In arm B the head's get at layer 0 is a hit. It re-plans the queue from the head's own one-member group, and that drops the planned read of layer 3. It happened in every policy forward (25 of 25 steps over 5 runs), plus once per run in step 0's reference forward, where the head is still a demand read: 30 drops in all, against 0 in arm A.
  • The cause is that the reference forward leaves layer 1 off the device, so layer 0's advance() issues get(1), which plans [2, 3], immediately before the head's get.
  • On this fixture the next decoder hit put layer 3 back in the queue in time. Steady-state demand misses were 46 in arm A and 42 in arm B, and the blocked-time median was 48 ms in both.
• SFT. The drop happens only at step 0, where the head is still a demand read ([3] at depth 3, [3, 4] at depth 4, 3 of 3 runs each). The steady state is identical.
• read_ahead=2. No drops in any mode. In the preference-shaped runs, steady-state misses were 82 in arm A and 83 in arm B.

None of this calls for a change to _plan_queue, _claim_slot or get, and I am not suggesting one. What it changes is the disk-tier run: it needs DPO (or KTO) at stream_read_ahead: 3 as well as SFT at the default depth, because that is the only configuration where the mechanism fires. Whether it costs any wait on a read-bound NVMe run is a question for the card.

Before this leaves draft

1. The measurement on the 5070, as set out on predicted_peak_bytes under-predicts VRAM for an untied streamed head once #1061 lands; is #975's one-slot ping-pong worth trading for 2 slots? #1174: fix(layer-stream): prefetch embed_tokens at the backward tail, not the next step's prime (#975) #1178's smaller fixture, judged by the step's total wait (every entry in stalls), on the RAM and the disk tier, plus the preference, read_ahead: 3 case above. The default of 0 should follow from that number. If layer 2's wait or the disk tier gets worse, head_prefetch_layer already accepts a later layer.
2. The docs sentence. docs/performance-and-quantization.md:356 still says the shared slot "is reused for the output head after the last decoder layer". With this default, the head's copy is issued right after the embedding lookup, so that sentence needs to say so.
3. A rebase once fix(layer-stream): prefetch embed_tokens at the backward tail, not the next step's prime (#975) #1178 lands. Today your branch merges cleanly with main (afd9a917) and with the other open drafts in layer_stream_runtime.py, feat(stream): prepare real preference VRAM probes (#840) #1141 and feat(stream): make NF4 fused forward checkpoint-visible (#842) #1143.
4. The fragment, after the measurement. It describes the layer-0 default as shipped behaviour, so it has to match whatever the number decides. Its credit line is already right.

Non-blocking

• The range check. It runs only in the constructor, and the harness switch is an attribute assignment, so the switch skips it. An out-of-range value then surfaces at the head's forward as large-layer scheduler bug: slot holds ..., which points at the wrong cause. A property that runs the same check would keep one rule. True also passes, as layer 1.
• What the bit-identity class covers. It passes 8 of 8 on fix(layer-stream): prefetch embed_tokens at the backward tail, not the next step's prime (#975) #1178's head as well, where the attribute does nothing, so it guards against regressions rather than pinning the move. The move itself is pinned by the trigger tests, test_the_head_load_comes_at_the_configured_layer and the installed-default test. That split is fine. It helps to know it when one of them goes red.

CI: all 13 required checks pass on d9fc2a17.

Not run here, because it needs a GPU:

• the pool-level GPU test in both directions (you report 5 of 5 and 5 of 5 on your 3050);
• the file under CUDA (you report 23 passed);
• the measurement itself.

[VanshSharmaPES]

Thanks for the thorough pass, and for comparing the two trees and not only the two arms. I've pushed 427befb for two of the items; the rest are waiting on the measurement or on #1178.

The range check. You're right that the harness switch is an assignment, so it skipped the constructor's check and a bad value would only surface at the head's forward as "scheduler bug". head_prefetch_layer is now a property whose setter runs the same check, and the constructor goes through it. It also refuses bool now, so True no longer passes as layer 1 (and False no longer as layer 0). A refused assignment leaves the old value. I added 13 test cases: two more constructor cases (True, 1.0), seven assignment cases that check the error and that the value is unchanged, and four valid assignments (None, 0, 1, 3). With the setter's check replaced by if False:, exactly the 12 refusal cases fail.

The docs sentence. docs/performance-and-quantization.md:356 now says the head's copy is issued right after the embedding lookup (StreamPrefetcher.head_prefetch_layer, default 0) instead of at the last decoder layer. It describes the timing only and makes no claim about speed, since that isn't measured.

Checks on 427befb: ruff is clean; on CPU pytest -k "stream or 1049 or 975 or 1174 or docs or changelog" gives 637 passed, 36 skipped, 0 failed; and on the 3050 (torch 2.14.0+cu130) both test files give 41 passed. Not re-run: the two GPU tests in both directions, since nothing they cover changed.

Left as is:

• The changelog fragment, because it has to match what the measurement decides about the default.
• The rebase, until fix(layer-stream): prefetch embed_tokens at the backward tail, not the next step's prime (#975) #1178 lands.
• The prime() layer-0 re-copy, which stays with its own issue.

Your note about the bit-identity class taking 8 of 8 on #1178's head is fair. I'll keep the split you described: those tests guard against regressions, and the trigger and installed-default tests pin the move itself.

I'm also planning to write the interleaved A/B script for the 5070 run, covering both tiers, total stall, and DPO at stream_read_ahead: 3, so it's one command. I'd only smoke-test it on the 3050 for correctness, and I won't report those numbers as a result.

[VanshSharmaPES]

Pushed b1c595c: an interleaved A/B driver, benchmarks/harness/head_prefetch_ab.py, so the 5070 run for this PR is one command.

It flips StreamPrefetcher.head_prefetch_layer between blocks in one process on one model, the way #1178 ran backward_tail_prefetch=None as arm A. It uses the committed stream_probe.Instruments and layer0_wait's labelling, so it's the same wait() brackets as gate-971 §8. It reports the step's total wait per arm plus per-group medians (embedding, head, layer 0, other decoder layers), and it covers both tiers, --read-ahead, and a --step-shape preference (a no-grad forward first, not a TRL loss) for the disk-tier case at --read-ahead 3, which is the only configuration you found the reader's queue gets replanned in. Default schedule is 3 warm-up steps, then 5 blocks of 1 discarded + 2 recorded steps per arm, with the arm order rotating each block.

On a tree without head_prefetch_layer it exits with a message instead of comparing an arm with itself — I checked that against a worktree of #1178's head.

I ran it on my 3050 on the RAM tier and on the disk tier with the preference step at read-ahead 3, with a 4-layer untied fixture I generated locally, only to confirm it completes and the switch takes effect. That's 4 samples per arm on a different card, so I'm not reporting those numbers as a result, and the synthetic generator has no hidden-2048 shape to match your fixture anyway. It should be ready to point at whatever fixture the 5070 run uses.

Example invocations are in the module docstring.

[MakazhanAlpamys]

@VanshSharmaPES, thank you for the driver. It turns the 5070 run into commands instead of a procedure. Three things about it stand out:

• It judges an arm by the step's total wait and prints the per-group medians beside it. That is the right number to judge by: the head's own entry can improve while a decoder entry behind it on the copy stream gets worse.
• The arm order rotates every block, and the JSON is written after each one, so a run that dies keeps its earlier points.
• It refuses to run on a tree that has no head_prefetch_layer, instead of comparing an arm with itself.

What I checked (on CPU; nothing here ran on a GPU)

• --help works. --arms last,last is refused. Without CUDA it prints SKIP: CUDA is required and writes no JSON.
• Every probe API it calls exists with the signature it uses. The label wiring (_label_wrappers, then zip(stall_labels, inst._stall_pairs)) matches layer0_wait.py's own main(). stream_probe.build() reads only attributes your Namespace sets, so the first GPU run will not die on a missing argument.
• The "no head_prefetch_layer" branch is only reachable with CUDA. I confirmed it by reading, not by running.
• --arms last,7 parses. Layer 7 is the last layer of fix(layer-stream): prefetch embed_tokens at the backward tail, not the next step's prime (#975) #1178's 8-layer fixture, so last,7 runs the same code under two labels. That makes it a free A/A noise-floor control.

One correction to the rebase plan

On #1178 you wrote that you would run git rebase --onto origin/main <this PR's squash commit>. That form conflicts. I tried it on a local squash of #1178's current head onto today's main. It stops at 3dcd7e70 with conflicts in layer_stream_runtime.py and tests/test_issue975_backward_tail_embed_prefetch.py. The squash commit is not in this branch's history, so git replays all six of #1178's original commits as well. The upstream to cut at is the commit this branch forked from:

git rebase --onto origin/main 77dc04e3

That form replays your five commits cleanly. Afterwards layer_stream_runtime.py, async_disk_source.py, the driver, stream_probe.py, layer0_wait.py and your test file are byte-identical to b1c595c9. On that tree, the #975 and #1174 test files plus the changelog-fragment tests give 61 passed, 2 skipped (the two GPU tests), 0 failed, and ruff is clean.

Before this is ready for review

1. The RTX 5070 measurement. It needs fix(layer-stream): prefetch embed_tokens at the backward tail, not the next step's prime (#975) #1178's fixture and that card, so nothing on your side blocks it. Measuring today's head measures the code that would land, as long as main does not change those files first.
2. The rebase above, once fix(layer-stream): prefetch embed_tokens at the backward tail, not the next step's prime (#975) #1178 lands.
3. The default and the changelog fragment, set to whatever the number decides.
4. A row for head_prefetch_ab.py in the harness table of benchmarks/README.md. Every other harness script has one there, stream_probe.py and layer0_wait.py included.
5. Worth doing before the run: record soup_cli.__file__ (and the git SHA if you can) in the JSON payload, the way stream_probe's meta has since Async NVMe disk source is 1.03-1.20x slower than sync DiskSource when the store is fully page-cached #974. A record should say which tree it measured. Today only --label can carry that. Optionally, add a --shards passthrough, so the disk tier can read the same NF4 store fix(layer-stream): prefetch embed_tokens at the backward tail, not the next step's prime (#975) #1178's A/B used instead of re-sharding.

CI: test (windows-latest, 3.12) failed on b1c595c9 with Windows fatal exception: code 0xc000001d, inside tests/test_issue1235_prm_fp32_master_weights.py. That is a runner-CPU fault in a test this PR does not touch. main fixed it in #1314, which merged after this run started. It clears on your next push. The other 12 required checks pass.

[MakazhanAlpamys]

@VanshSharmaPES, a quick follow-up. #1178 merged as 20ab4668 a few seconds before my last comment went up, so the rebase in step 2 of that list ("once #1178 lands") is no longer waiting on anything.

I tried the command against the real merge commit, on today's main (df106781):

git rebase --onto origin/main 77dc04e3

It replays your five commits with no conflicts. layer_stream_runtime.py, the driver, your test file and the fragment come out byte-identical to b1c595c9, and your docs hunk is unchanged. On that tree, on CPU with the GPU hidden, the #975 and #1174 test files plus the changelog-fragment tests give 61 passed, 2 skipped, 0 failed. The push also brings in #1314, which should clear the test (windows-latest, 3.12) failure.

What is left on your side, all from my last comment:

• the rebase above;
• a row for head_prefetch_ab.py in the harness table of benchmarks/README.md;
• soup_cli.__file__ (and the git SHA if you can) in the JSON payload. The --shards passthrough stays optional.

They can all go in one push. The default and the fragment still follow the RTX 5070 number. That run is on my side, and nothing above waits for it. If any of this is in your way, or you are busy this week, just say so here.

[MakazhanAlpamys]

@VanshSharmaPES, thank you for doing all three items in one push. Three things about it stand out:

• The rebase is exact. git range-diff shows all five commits replayed as identical patches (5 of 5 =). layer_stream_runtime.py, your test file and the fragment are byte-identical to b1c595c9, so the RTX 5070 run below measured exactly the code we reviewed.
• The setter holds up. I broke it eight different ways, and every one turns your file red. Your own if False: mutation fails exactly the 12 refusal cases, as you said.
• The driver worked on the card. It took one command per configuration, and the last,7 A/A arms gave the noise floor the decision needed.

The RTX 5070 result

The run used your new head merged with today's main (bb92aa2c) on #1178's 8-layer untied fixture: hidden 2048, vocabulary 128,256, a 501 MiB shared slot, NF4, sequence 512, batch 1. It kept the driver's default schedule, which gives 10 recorded steps per arm. Median total wait per step, in ms:

run	last	0	change	A/A spread
RAM tier, SFT	11.96	4.50	−7.46	0.11
disk tier, SFT, read-ahead 2	15.83	1.42	−14.41	0.74
disk tier, preference step, read-ahead 3	57.29	29.16	−28.13	—

• Every sample separates. In all three runs, every layer-0 sample is below every last-layer sample.
• The head's wait is gone, and the decoder layers pay part of it back. The head's own wait drops to 0.1–0.2 ms. As you predicted, the other decoder layers now wait longer, because a decoder copy queues behind the head's on the shared copy stream: 0.78 → 4.06 ms on the RAM tier and 1.03 → 9.24 ms in the preference step. On disk SFT there is no change (0.91 → 0.93). That gives back part of the head's saving, not all of it.
• Layer 0 is the best place, so the default stays at 0. A RAM-tier sweep of last,0,1,2,3 gave 12.49 / 3.99 / 4.80 / 4.97 / 5.06 ms. The sweep ran on the pre-rebase head, whose layer_stream_runtime.py is byte-identical.
• The GPU guard is pinned on the card. Both pool GPU tests pass on the new head, 3 of 3 runs. With the wait_stream removed from LargeLayerBufferPool.load_async, both fail, 3 of 3.

Step time moved the same way (medians 230.0 → 222.9, 294.4 → 286.9 and 397.5 → 381.3 ms). Step time is noisy on this box, though, so the wait is the number to quote.

Before this leaves draft

1. The fragment. changelog.d/0.75.1/1258.changed.md still says "The speed effect is not measured yet: it needs the step's total wait() on a GPU, on both the RAM and the disk tier." It can now say what was measured. A suggested replacement for that sentence:

On an RTX 5070 Laptop GPU, with an 8-layer untied checkpoint (hidden 2048, a 501 MiB shared slot, NF4, sequence 512), the step's median total wait() fell from 11.96 to 4.50 ms on the RAM tier, from 15.83 to 1.42 ms on the disk tier, and from 57.29 to 29.16 ms on the disk tier at stream_read_ahead: 3 for a step that runs a no-grad forward first. A decoder layer's copy now queues behind the head's on the shared copy stream, so the other decoder layers wait longer (0.78 to 4.06 ms on the RAM tier), which gives back part of the saving; layer 0 was the best of layers 0 to 3.

2. The NVIDIA driver version is lost in every new JSON (benchmarks/harness/head_prefetch_ab.py:297-299).

• The cause. stream_probe.gpu_facts() returns its own "driver" key, which holds the NVIDIA version. The payload spreads it first and then sets "driver": "benchmarks/harness/head_prefetch_ab.py", and the later key wins.
• The evidence. In the five JSONs from the card, 616.92 appears nowhere. In the seven JSONs from b1c595c9 it is at gpu.driver.
• A second effect. gpu also changed from a dict to the bare device name, so old and new records no longer have the same shape.

One correction on my side, because it led to this. On 10:16Z I wrote that only --label could say which tree a run measured. That was wrong. gpu_facts() has recorded soup_cli_file since #991 (the #974 work), and your b1c595c9 payload already had it under gpu. So the smallest fix is to put the nesting back and keep your new key beside it:

        "driver": "benchmarks/harness/head_prefetch_ab.py",
        # the card, its NVIDIA driver and PCIe link, torch, and soup_cli_file
        "gpu": stream_probe.gpu_facts(device),
        "git_sha": _source_sha(),

Those two items are all I would hold this on.

Worth doing in the same push

3. git_sha names the driver's checkout, not the tree that was measured (head_prefetch_ab.py:106-129).

• What happens. _source_sha() runs git in Path(__file__).parents[2]. With PYTHONPATH pointing at another tree, it recorded 24ba7690 (the driver's checkout) while soup_cli was imported from bb92aa2c.
• A precedent in the repo. benchmarks/harness/variant2_gate.py:213-256 resolves the commit from soup_cli.__file__ with git rev-parse --show-toplevel and appends -dirty, which also covers a run on a tree with an uncommitted change. unknown for a tree with no .git, as on the card, is honest and fine.
• A test. This one is in the style of tests/test_issue379_variant2_gate.py::TestSourceCommit. It fails 3 of 3 on 77365646, and passes 3 of 3 with the nesting fix above plus a _source_sha() written the way variant2_gate.py writes it:

"""#1174: the head-prefetch A/B driver's JSON must say which card, driver and tree it measured."""

import importlib.util
import shutil
import subprocess
from pathlib import Path
from types import SimpleNamespace

import pytest

_HARNESS = Path(__file__).resolve().parents[1] / "benchmarks" / "harness" / "head_prefetch_ab.py"


def _load_driver():
    spec = importlib.util.spec_from_file_location("head_prefetch_ab", _HARNESS)
    assert spec is not None and spec.loader is not None
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


_driver = _load_driver()


@pytest.mark.parametrize("status", ["", " M src/soup_cli/utils/layer_stream_runtime.py\n"])
def test_git_sha_names_the_imported_tree_and_marks_it_dirty(monkeypatch, tmp_path, status):
    """The driver and the soup_cli it measures can come from different checkouts, so
    the commit is resolved from ``soup_cli.__file__``, not from this file."""
    import soup_cli

    package_dir = Path(soup_cli.__file__).resolve().parent
    calls = []

    def fake_run(command, *, cwd, **kwargs):
        calls.append((list(command[1:]), Path(cwd)))
        if command[1:] == ["rev-parse", "--show-toplevel"]:
            return SimpleNamespace(stdout=f"{tmp_path}\n", returncode=0)
        if command[1:] == ["rev-parse", "HEAD"]:
            return SimpleNamespace(stdout="a" * 40 + "\n", returncode=0)
        return SimpleNamespace(stdout=status, returncode=0)

    monkeypatch.setattr(shutil, "which", lambda name: "git")
    monkeypatch.setattr(subprocess, "run", fake_run)

    assert _driver._source_sha() == "a" * 40 + ("-dirty" if status else "")
    assert calls[0] == (["rev-parse", "--show-toplevel"], package_dir)


def test_the_gpu_facts_stay_nested_so_the_nvidia_driver_survives():
    """``gpu_facts`` returns a ``driver`` key (the NVIDIA version); spreading it into a
    payload that also sets ``"driver"`` to the script path overwrites it."""
    text = _HARNESS.read_text(encoding="utf-8")
    assert '"gpu": stream_probe.gpu_facts(device)' in text
    assert "**stream_probe.gpu_facts(" not in text

4. Warm the shard cache first, and keep paths short on Windows. The card run hit two things that belong in the driver's docstring:

• Sharding in the same process slows both arms. The first process that also built the NF4 shards measured both arms about twice as slow as later processes: the last-layer head wait was 27.2 ms, against 10.9–11.5 ms in the five later RAM-tier processes. Its change (−16.45 ms) therefore overstated the effect; the three later RAM-tier A/B runs agree at −7.46 to −8.49 ms. The docstring should say to build the shards once and discard that run.
• A deep weights path fails on Windows. It stops with WinError 206 (MAX_PATH), because resolve_shard_dir names the shard directory after the weights path, up to 128 characters. That is layer_shard.py's behaviour, not your driver's.
• One change fixes both. The --shards passthrough from my earlier list now covers both cases. stream_probe.build() already reads args.shards, and the driver hard-codes shards=None (:208).

Nits (optional)

• ruff format would rewrite head_prefetch_ab.py:333-334 and the setter's condition (layer_stream_runtime.py:1027-1031). CI runs only ruff check, and that is clean.
• test_an_assignment_is_checked_like_the_constructor_and_leaves_the_value starts from 0. For False, the "value unchanged" check therefore cannot tell a stored False from the old 0, and that one case survives a store-before-check mutation. Starting from 2 closes the gap.

What I ran

Everything here ran on CPU with the GPU hidden, on your head merged with today's main. That merged tree is the tree the card measured (48a012ac).

• Tests. The #975 and #1174 files, the changelog-fragment and merge-freshness tests, and the two benchmarks/README.md index tests: 104 passed, 2 skipped (the two GPU tests), 0 failed. ruff check is clean, and all 185 fragments validate.

• Your file on main without your branch. 24 failed, 11 passed, 1 skipped. The 11 that pass are the regression guards we discussed before.

• Eight mutations of the setter, all killed:

  mutation	result
  the whole check off (if False)	12 failed
  bool refusal dropped	3 failed
  int check dropped	3 failed
  upper bound <=	2 failed
  lower bound dropped	2 failed
  value stored before the check	6 failed
  constructor bypasses the property	5 failed
  getter always None	15 failed

• Three mutations of the provenance. I dropped gpu_facts, hard-coded git_sha and made _source_sha() always return unknown. None of them fails a test, because no test runs the driver yet. The test in item 3 is for that.

• Merges. Your branch merges cleanly with #1141 and #1143, and with the other open PRs that add rows to benchmarks/README.md.

CI: all 13 required checks on 77365646 are still queued.

Once items 1 and 2 are in and the checks pass on that head, nothing else on my list stands between this and review.

[VanshSharmaPES]

Thanks for catching both of those, and for the correction on where the tree-provenance fix should live — I hadn't noticed the spread was clobbering gpu_facts()'s own driver key.

The lost NVIDIA driver. Fixed as you suggested: gpu_facts() is nested back under "gpu", and git_sha sits beside it. gpu is a dict again, so old and new records have the same shape.

git_sha named the wrong tree. Rewrote _source_sha() the way variant2_gate.py's is — resolves --show-toplevel and HEAD from soup_cli.file's directory, not from this driver's own file, and appends -dirty off git status --porcelain --untracked-files=normal. unknown on a tree with no .git, same as before.

Test. Added tests/test_issue1174_head_prefetch_ab_provenance.py, close to your sketch: the parametrized dirty/clean case against a faked subprocess.run, plus two unknown-path cases (no toplevel, no git binary), plus a source-text check that "gpu": stream_probe.gpu_facts(device) is present and **stream_probe.gpu_facts( is not, so a future spread-back regression fails loudly instead of silently.

--shards. Added, wired through to stream_probe.build(). Docstring now says to shard once, discard that process's numbers, and point --shards at the sharded directory afterward — covers both the ~2x slower first process and the Windows MAX_PATH case.

The fragment. Took your suggested replacement for the "not measured yet" sentence verbatim.

Checks on c7cafea: ruff check clean; on CPU, pytest -k "stream or 1049 or 975 or 1174 or docs or changelog" gives 677 passed, 13 skipped, 0 failed, plus the new provenance file at 5 passed on its own.

Left as nits, per your list: the ruff format rewrite and the setter test's starting value (0 → 2). Happy to take those too if you'd rather not carry them.

[VanshSharmaPES]

Picked up both nits in c14e45c, on top of c7cafea:

ruff format. Applied to the setter's condition (layer_stream_runtime.py:1027-1031) and the driver's one long print (head_prefetch_ab.py). I left the unrelated StreamedLargeLayer.init reformat ruff format also wants — that one's outside what you flagged.

The setter test's starting value. test_an_assignment_is_checked_like_the_constructor_and_leaves_the_value now starts from head_prefetch_layer=2 instead of 0, so a store-before-check mutation storing bad=False can't hide behind Python's False == 0 — the "value unchanged" assertion now actually distinguishes them.

Checks on c14e45c: ruff check clean, ruff format --diff clean on every file this PR touches (the StreamedLargeLayer.init hunk aside), and on CPU pytest -k "stream or 1049 or 975 or 1174 or docs or changelog" gives 677 passed, 13 skipped, 0 failed.

[MakazhanAlpamys]

@VanshSharmaPES, thank you. All four items are in c7cafeaa, pushed within minutes of the review, and both optional nits are in c14e45cf. Your commit message says plainly what broke and why, which makes a review like this quick. Three things stand out:

• The new test file pins both bugs. Put the top-level spread back, or resolve the commit from the driver's own file again, and it fails. It also covers -dirty and both unknown paths.
• The fragment is exact. I recomputed every number in it from the raw JSONs of the run: 11.96 → 4.50, 15.83 → 1.42, 57.29 → 29.16, the decoder rise 0.78 → 4.06, the 501 MiB slot, 8 layers, sequence 512, and layer 0 as the best of 0–3. Each one matches.
• The nits are real fixes, not just style. The setter's reformat changes no code (the AST of StreamPrefetcher and of the driver is identical before and after). Starting the assignment test from 2 means a store-before-check mutation now fails 7 cases instead of 6, False included.

What I ran

Everything ran on CPU with the GPU hidden, on c14e45cf merged with today's main (bb92aa2c). After main moved to f76c8afd, I ran the test set again and got the same counts.

Tests. The #975 and #1174 files, your new provenance file, the changelog-fragment and merge-freshness tests, and the two benchmarks/README.md index tests: 109 passed, 2 skipped (the two GPU tests), 0 failed. ruff check is clean. ruff format --check is clean on the driver and both test files; the one hunk left in layer_stream_runtime.py is main's own.

The driver's real main(). The payload is built inside main() behind the CUDA check, so I faked only the GPU-bound pieces: stream_probe, layer0_wait, the bitsandbytes optimizer, and the three torch calls made with device="cuda". I faked gpu_facts with its real key set; everything else ran as your code.

• On c14e45cf, gpu is a dict and gpu.driver is 616.92. git_sha is the merged tree's commit, and --shards D:/fake/untied-8l-shards reaches build() and the JSON's shard_dir.
• The same run on 77365646 gives gpu as a bare string, with no driver version anywhere.

_source_sha() with real git.

• A git archive extract (no .git) gives unknown.
• The clean worktree gives its commit.
• A tracked edit under src/ gives -dirty.

Eight mutations of the setter, all killed (12, 3, 3, 2, 2, 7, 5 and 15 failures). The store-before-check one is the 7, which was 6 before your change.

Nine mutations of the provenance code, five killed:

mutation	result
gpu_facts spread at the top level again	1 failed
gpu_facts dropped	1 failed
_source_sha() never reports the commit	2 failed
commit resolved from the driver's own file again	2 failed
-dirty dropped	1 failed

Two of those five survived at 77365646: dropping gpu_facts, and _source_sha() never reporting the commit. The four that still pass are listed below, and none of them is something I would hold this on.

From my side this is ready. Nothing on my list is left, so you can take it out of draft. It then waits only on the 13 required checks, which are queued on c14e45cf right now.

Optional follow-ups (none of them holds the merge)

1. Two sentences in the docstring (head_prefetch_ab.py:31-39).

• Who measured it. "I saw a last-layer head wait of 27.2 ms…" reads as your measurement, but those numbers came from the RTX 5070 run on this PR. Since benchmarks/ is the tracked record, it is worth saying so.
• What caused it. "The sharding itself runs on the same card as the arms, so a process that also shards measures both arms slower" states a cause. NF4 quantisation does run on the card (layer_shard.py:905-924), but it finishes before the first step, and why the arms were slower afterwards was not isolated.

Something like:

The RTX 5070 run on #1258 measured a last-layer head wait of 27.2 ms in the process that also built the NF4 shards, against 10.9-11.5 ms in five later processes; the cause was not isolated. Build the shards once and discard that process's numbers, or point --shards at an already-sharded directory.

2. Two error branches have no test. Dropping the 40-hex check changes nothing your tests can see, and neither does narrowing except (OSError, subprocess.SubprocessError) to OSError.

The second one matters in practice. The card run used a git archive tree, where git rev-parse --show-toplevel exits 128, and with check=True that is a CalledProcessError.

The two tests below pass on c14e45cf, and each fails one of those mutations. They are ready to append to your provenance file. They use its _driver, SimpleNamespace, Any and Path, plus import subprocess:

def test_source_sha_is_unknown_when_git_exits_non_zero(monkeypatch: pytest.MonkeyPatch) -> None:
    """A tree with no ``.git`` (a ``git archive`` extract) makes ``git rev-parse
    --show-toplevel`` exit 128; with ``check=True`` that is a ``CalledProcessError``."""

    def fake_run(command: list[str], *, cwd: Path, **kwargs: object) -> Any:
        raise subprocess.CalledProcessError(128, command)

    monkeypatch.setattr(_driver.subprocess, "run", fake_run)
    assert _driver._source_sha() == "unknown"


def test_source_sha_does_not_record_a_head_that_is_not_a_commit(
    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
    def fake_run(command: list[str], *, cwd: Path, **kwargs: object) -> Any:
        if command == ["git", "rev-parse", "--show-toplevel"]:
            return SimpleNamespace(stdout=f"{tmp_path}\n")
        if command == ["git", "rev-parse", "HEAD"]:
            return SimpleNamespace(stdout="not-a-commit\n")
        return SimpleNamespace(stdout="")

    monkeypatch.setattr(_driver.subprocess, "run", fake_run)
    assert _driver._source_sha() == "unknown"

3. The payload wiring itself is still untested. The other two survivors are hard-coding "git_sha": "unknown" in the payload and reverting shards=cli.shards to None. Both live in main() behind the CUDA check. A CPU test that runs main() with the GPU-bound pieces faked would cover them, the way tests/test_issue379_variant2_gate.py already runs its harness's main().

4. The record does not say how it was configured. The JSON has tier, read-ahead, sequence, batch and step shape. It has no quant, lora_r, lora_targets, buffers, seeds or no_pin, so the fragment's "NF4" cannot be read from the record itself. Adding "args": vars(cli) to the payload would capture all of them.

5. -dirty also fires on an untracked output file. With --untracked-files=normal, an untracked ab_ram.json at the repo root turns a clean tree into …-dirty. The docstring's --out ab_ram.json, run from the root, does exactly that from the second run on. It errs in the safe direction, and variant2_gate.py behaves the same way. Adding -- src to the status call, or writing outputs outside the tree, would avoid it.

[VanshSharmaPES]

Picked up all five optional follow-ups too, in d7f5551 on top of c14e45c:

1. The docstring. Reworded to say the 27.2 ms number came from the RTX 5070 run on this PR specifically, and dropped the cause I hadn't isolated (that sharding itself slows the arms).

2. -dirty scoping. _source_sha()'s git status now runs -- src, so an untracked --out file left at the repo root doesn't turn a clean tree dirty.

3. Two more _source_sha() tests. Added the non-zero git-exit case (a git archive tree — what the card's run actually had) and a HEAD that isn't a 40-hex commit. Both fail without the fix.

4. The payload wiring itself. Added a CPU test that runs the driver's real main() with only the GPU-bound pieces faked (model, runtime, stream_probe, layer0_wait, the bitsandbytes optimizer), then asserts on the JSON it actually writes. A hard-coded git_sha or shards reverted to None now fails a real run of main(), not just a source-text check.

5. "args" in the payload. Added "args": vars(cli), so a record says how it was configured (quant, lora_r, seeds, --shards, etc.) instead of relying on the changelog fragment's word for it.

Checks on d7f5551: ruff check clean, ruff format --diff clean on every file this PR touches, and on CPU pytest -k "stream or 1049 or 975 or 1174 or docs or changelog" gives 680 passed, 13 skipped, 0 failed (8 of those in the extended provenance file alone).

[MakazhanAlpamys]

@VanshSharmaPES, thank you for taking all five. The main() test is the one I'd point others at: it runs the driver's real code path on CPU, so the JSON is now checked as it is written, not as the source reads.

I verified each of the five by running it on d7f55514 merged with today's main (f76c8afd), on CPU with the GPU hidden:

1. The docstring now credits the RTX 5070 run on this PR and says the cause was not isolated. The first-person "I saw" and the stated cause are gone. The only code changes in the driver are the -- src scope and the "args" key; everything else in the diff is docstring or comment text.
2. The two _source_sha() tests pass. Without the 40-hex check, and with the except narrowed to OSError, one of them fails in each case.
3. The main() test passes. Hard-coding "git_sha" in the payload, reverting shards=cli.shards, and dropping "args" each fail it.
4. "args": a real main() run writes all 21 flags, including quant: nf4, lora_r, lora_targets, buffers, both seeds, no_pin and shards.
5. -- src, checked with real git:
   • an untracked ab_ram.json at the repo root no longer marks the tree dirty;
   • an untracked file or a tracked edit under src/ still does;
   • a git archive tree still gives unknown.

Tests: the #975 and #1174 files, your provenance file (8 of 8), the changelog-fragment and merge-freshness tests, and the two benchmarks/README.md index tests give 112 passed, 2 skipped (the GPU tests), 0 failed. ruff check and ruff format --check are clean.

Mutations: all 9 of the provenance mutations from my last review fail your tests now; four of them passed before this commit. Two more that I added also fail: dropping "args" and unscoping the git status call. That is 11 of 11.

A footnote on my own suggestion, and no change is needed here: -- src also stops a local edit to benchmarks/harness/*.py from marking the tree dirty. When the driver and soup_cli come from the same checkout, -- src benchmarks/harness would cover the measuring scripts too.

Nothing is open on my side. The PR is still a draft; once it is marked ready, it waits only on the 13 required checks, which are queued on d7f55514.

[MakazhanAlpamys]

Merged as 15a66781. Thank you, @VanshSharmaPES: you built the instrument before there was a number. The timing switch runs both arms in one process, and the A/B driver you wrote for a card you do not have ran on it as one command per configuration and produced a clear result. When review found that the provenance change dropped the NVIDIA driver version from every record and named the wrong tree, the fix and tests pinning both came within ten minutes. You then also took all five optional follow-ups, including a CPU test that runs the driver's real main(). On the tree merged with main, 112 tests passed and the 2 GPU tests skip on CPU. All 8 mutations of the new validation and all 11 of the provenance code fail the tests. On the RTX 5070 both pool GPU tests pass 3 of 3, and fail 3 of 3 with the wait_stream removed.