[Amix29]

Refs #842
Refs #841

What

This PR implements checkpoint-visible streamed NF4 and folds the #841 pooled-view cache into the same branch. It follows the maintainer ruling that native bitsandbytes is the numerical reference.

• reads the bitsandbytes private dispatch cap, alignment rule, and heuristic at call time;
• uses the Soup checkpoint-visible Function only when bitsandbytes selects its custom CUDA GEMM;
• keeps dequantize_4bit plus F.linear when bitsandbytes selects that fallback;
• fails safely toward the Function if the private dispatch symbols move;
• saves packed data and quant-state tensors through save_for_backward, restores autocast state, skips unneeded gradient work, and asserts stream-slot ownership before backward reads pooled aliases;
• caches Params4bit and QuantState substitution views by pool mapping identity without retaining another weight copy;
• rewrites Pin the bitsandbytes fused-vs-dequant NF4 divergence that #857 rests on #968 to require patched resident == unpatched native resident at a fused M.

Validation after rebasing on #1258

• six focused compatibility files: 166 passed, 23 expected platform/GPU skips;
• full Ruff and git diff --check clean.

CUDA boundary

On the pre-rebase head, the maintainer's RTX 5070 G1 suite passed 19/19 and all three G2 mutations were killed. Because the requested rebase on #1258 rewrote layer_stream_runtime.py, those same G1 and G2 gates must be rerun on the current head before merge. No CUDA result is claimed from this M4 Max.

The backward still dequantises because bitsandbytes exposes no deterministic transposed 4-bit GEMM, so the full no-dense-W ceiling remains open at issue level.

[MakazhanAlpamys]

Advisory review of the draft at c326a396. Keep it in draft. The mechanism is sound, but the largest effect is missing from the description.

What actually changes. On CUDA, install_dequant_forward no longer dequantises: its forward now runs bitsandbytes' own gemm_4bit dispatch, the same one native MatMul4Bit runs. So the streamed NF4 reference moves from "dequant + F.linear" (#331 / #857) to "native bitsandbytes". #842's own "Honest limits" section says that change should be decided deliberately, not arrive as a side effect. The PR also deletes #857's warning ("not assumed to be numerically free") at the moment it makes the switch.

Measured on an RTX 5070 with a 4-layer / 2-buffer streamed model, bit for bit:

• With the PR: streamed == native resident NF4, for the logits and every LoRA gradient, in bf16 and fp16, at seq 12 and 128.
• On main: streamed != native at seq 12 (logits 4.883e-3 / grads 1.826e-3 in bf16).

So the #776 symptom disappears because streaming moves onto bnb's native dispatch, not because anyone settled which path is right. I will give a ruling on that choice separately. Please don't build further on either answer until then.

1. Blocking: two CUDA tests go red. test_issue968::TestTheTwoBitsandbytesFourBitPathsDisagree::test_the_paths_diverge_by_the_measured_magnitude[float16|bfloat16] fails with "the two bitsandbytes 4-bit paths now AGREE … assert 0.0 > 0.0", on this head and on the tree merged with main. Clean main gives 167 passed / 0 failed on the same files. #968's patched arm is resident + install_dequant_forward, which is now the same dispatch as native. CI cannot see this, because the runners have no GPU. There are two ways out:

• Accept the new reference. Then #968 needs rewriting, and these still say the streamed path always dequantises: test_v07202.py:1916, test_v07300.py:35 ("WHY THIS IS NOT A NUMERICS CHANGE"), :233 and :296.
• Keep install_dequant_forward's arithmetic, and put the Function on its own path.

2. Blocking: the tests do not pin the CUDA change. All 6 new tests are CPU-only. On CPU, gemm_4bit is bnb's default dequantize_4bit + F.linear, so "fused vs dequant exact parity" compares the old arithmetic with itself. At 16x16, bnb falls back even on CUDA (K % 64 != 0). Re-run on CUDA at fused shapes (64x64 and 4096x4096, M=3), parity is False, with a bf16 forward gap up to 1.25e-1. Three mutations survive:

• Gating the new route off entirely: the CUDA files get greener (173 passed / 0 failed).
• Recomputing with dequant + F.linear during checkpoint recompute: a gradient probe shows LoRA grads off by 1.28e-3 (bf16, seq 12), and nothing in the repo notices.
• Routing the CUDA branch back to bnb.matmul_4bit, the #331 route itself: TestStreamedNF4AvoidsMatMul4Bit builds a CPU model, where that branch is unreachable. A counter on a CUDA streamed step reads matmul_4bit = 40 under this mutation, against 0 on the PR.

Needed as gpu-marked tests:

• (i) the matmul_4bit counter on a CUDA streamed model, plus a counter on the Function path;
• (ii) recycled-vs-private gradient comparison through the same Function at a fused shape (K a multiple of 64, M inside the fused window), with the plain-ctx arm that must fail;
• (iii) streamed-vs-resident logits and LoRA grads at a fused M.

3. Blocking: it is slower at training shapes, on the RAM tier #842 targets.

Dequantisations per checkpointed forward + recompute + backward of one 8B q projection (bf16):

• 3 on the PR vs 2 on main at M=512 and 2048;
• 1 vs 2 at M=16 and 64.

The fused window on this card: for 8B, q / o / down up to M=96, k / v up to M=8, gate / up up to M=256; for 70B, every projection up to M=256.

Timing one checkpointed 7-projection decoder region (indicative only, the card was shared at the time), new/old:

• 8B: 1.16-1.22 at M=16, 0.88-0.94 at 64, 1.03-1.19 from 128 to 2048;
• 70B: 0.82 / 0.88 / 1.02 / 1.08 at M=16 / 64 / 256 / 512.

Peak VRAM for the same region goes down: 8B by 260-301 MiB, 70B by 1.08-1.18 GiB. The Function saves packed aliases, where the old path held every dense W from recompute to backward; that is plausibly most of #331's +650.9 MiB. So this trades time for VRAM. It needs #842's gate, or it should route through the Function only where bnb picks the fused kernel, with the trade stated.

What holds:

• New == native bnb bit for bit at every shape, M and dtype, and dX is identical across new, old and native.
• The fused kernel is deterministic.
• Recompute and the original forward take the same bnb path, because the dispatch is a cached pure function of device, dtype and shape. So the shipped code has no "gradient for a different function" defect.
• The #841 per-layer view cache is correct and pinned (its test fails on main at assert second is first), and bf16 is now covered.

A new invariant worth writing down. The Function saves pool aliases and dequantises them in backward. Gradients are right only because checkpoint recompute calls pool.wait() first, and the prefetcher is 1-ahead with n >= 2, so no slot is refilled between a layer's recompute and the end of its backward. A deeper lookahead would break this silently. Add a comment, and an ownership assertion in backward.

Once it leaves draft it needs a changed fragment and docs. docs/performance-and-quantization.md:343 says streamed NF4 is bit-exact vs resident "through the same bitsandbytes kernels"; it must say which path CUDA streaming takes, and state the VRAM/time trade. None of #842's acceptance criteria are met yet, as the draft says.

Nits:

• backward computes grad_x and grad_bias regardless of ctx.needs_input_grad.
• The gate does an op lookup on every call.
• The Function is not autocast-aware (custom_fwd).

[MakazhanAlpamys]

The ruling I said I would give separately.

The reference for streamed NF4 is native bitsandbytes gemm_4bit. Streaming must match regular (resident) training bit for bit, on any card. So the direction of this PR is right: the #776 symptom going away is now the intended result, not a side effect. It also lifts the #776 caveat in #842's acceptance criteria. The streamed-vs-resident check has to hold on the RTX 5070 (sm_120) too.

It does not merge until these three are done.

(a) Remove the slowdown at training sizes. At training M the Function adds a third dequantisation per checkpointed step (3 vs 2 at M=512 and 2048), and my 8B region ran 3-19% slower from M=128 up (indicative, on a shared card). Route through the Function only where bitsandbytes itself picks the fused kernel. Everywhere else bitsandbytes runs dequantize_4bit + F.linear, which is exactly the old path, so both routes stay bit-exact with native.

• Take that decision from bitsandbytes' own dispatch, not from the window I measured, which belongs to this card. In 0.50.2 that is the _gemm_4bit_custom_max_m cap, the K % blocksize check and _gemm_4bit_use_custom_fn in bitsandbytes/backends/cuda/ops.py.
• Those names are private. If they are missing, fall back to the Function (always bit-exact, only slower), and add a test that fails when they move.
• If you also want the Function's VRAM saving at training sizes (260-301 MiB on 8B, 1.08-1.18 GiB on 70B in my region), put that route behind a flag, off by default, and state the time/VRAM trade in its help text.

(b) CUDA tests. Mark them gpu. CI skips them, so they get run on the RTX 5070 before merge.

• The matmul_4bit counter on a CUDA streamed step (expected 0), plus counters showing that the Function runs at a fused M and the old path at a training M.
• Recycled-vs-private gradients through the Function at a fused shape (K a multiple of 64, M inside bitsandbytes' window), with the plain-ctx arm that must fail.
• Streamed vs resident logits and every LoRA gradient, at a fused M and at a training M, against an unpatched resident model. test_issue385::TestFloat16StreamingIsBitExact applies install_dequant_forward to its resident arm (test: match NF4 resident reference compute path #857). Under this ruling the reference is plain bitsandbytes.

Between them these must kill the three mutations from my review: the counter catches the matmul_4bit route, and the unpatched reference catches the other two.

(c) Rewrite the #968 tests against the new reference. TestTheTwoBitsandbytesFourBitPathsDisagree::test_the_paths_diverge_by_the_measured_magnitude should pin that the patched and unpatched resident models now agree bit for bit at a fused M, with a counter proving the fused kernel ran. Four comments still say the streamed path always dequantises: test_v07202.py:1916 and test_v07300.py:35, :233 and :296. They should say what it does now.

The rest of my review still stands: the slot-ownership assertion in backward, a changed fragment, and docs/performance-and-quantization.md:343 naming the path CUDA streaming takes.

[Amix29]

Implemented the maintainer ruling and remaining review items at 8158f84c, rebased onto current main.

The Function now runs only when bitsandbytes itself selects the custom GEMM; training-M fallback stays on the old dequant+linear path. Missing private dispatch symbols fail safely toward the Function. Slot ownership, autocast decorators, needs_input_grad, docs, changed fragment, #968 rewrite, and the requested GPU counter/alias/streamed-vs-unpatched-resident matrix are included.

Local non-GPU validation: 133 passed, 11 skipped plus the rewritten #968 GPU module collected as 8 skipped; Ruff and diff-check clean. The PR stays draft until the RTX 5070 acceptance run is performed.

[MakazhanAlpamys]

@Amix29 thanks for turning the ruling around this quickly. The routing is what the ruling asked for. Here is what I checked at 8158f84c, and what the ruling means for this PR now. I ran everything on CPU: Windows, torch 2.14.0 with CUDA hidden, bitsandbytes 0.50.2.

What is good, by execution

• The gate follows bitsandbytes' own dispatch. _can_use_checkpoint_visible_nf4_gemm (layer_stream_runtime.py:335) makes bnb's own decision, the one in backends/cuda/ops.py:934-970: the same M/N/K, the _gemm_4bit_custom_max_m cap, the K % blocksize check and _gemm_4bit_use_custom_fn, all read at call time.

  • I drove bnb's registered CUDA kernel function itself on CPU: its two arms were stubbed, and the only thing faked was the device properties, set to sm_120. Your gate agrees with it in 2244 of 2244 cases (8B/70B projections and the fixture's shapes, M 1-2048, bf16/fp16/fp32, nested or not, 2-D and 3-D input). bnb picks its fused kernel in 944 of them.
  • With 36 SMs (the 5070 Laptop), bnb's heuristic reproduces exactly the 8B window I measured on the card: q/o/down fused up to M=96, k/v up to M=8, gate/up up to M=256.
  • If _gemm_4bit_use_custom_fn, _gemm_4bit_custom_max_m or the whole ops module is missing, the gate routes to the Function at both M=3 and M=512. That is the safe direction the ruling asked for.

• The Function works inside the real streaming runtime. On CPU, bnb's gemm_4bit is dequantise + linear, so I forced the gate on and trained 5 AdamW steps.

  • 2 layers / 2 buffers: 140 Function calls and 55 slot-ownership checks in backward, none of which fired.
  • 4 layers / 2 buffers, where slots are recycled: 280 calls and 125 checks, none of which fired.
  • In both setups, the forced run, the normal run and an unpatched resident model agree bit for bit: every loss, the step-0 logits and every LoRA gradient.
  • LoRA sits on q/v only, so k/o/gate/up/down ran the Function with no adapter next to them.

  This checks how it composes with the pool, recompute, the Layer streaming: cut the ~10 ms fixed cost per layer visit, which dominates the step below ~128 tokens #841 view cache and the ownership check. It does not check the CUDA kernel's numerics.

• The two correctness pins hold. Removing the ownership check is caught by test_backward_refuses_a_recycled_stream_slot. Moving the packed state from save_for_backward onto ctx (the Layer streaming: NF4 above ~165 MiB/layer produces silently wrong gradients (forward stays bit-exact) #331 class) is caught by the recycling test.

• (c) is done. The Pin the bitsandbytes fused-vs-dequant NF4 divergence that #857 rests on #968 rewrite (patched == unpatched at the fused M, with a counter proving the Function ran) is what I asked for, and all four stale comments are gone. A grep for "always dequantis", "NOT A NUMERICS CHANGE" and "not assumed to be numerically free" over tests/, src/ and docs/ finds nothing.

Tests I ran

• test_issue842, test_v07300 and test_v07202: 135 passed, 9 skipped, the same 144 tests as your 133 / 11.
• test_issue968: 8 GPU tests collected and skipped, as you reported.
• On the tree merged with main, the same files plus test_issue385 and test_issue487: 177 passed, 22 skipped, 1 failed. The failure is item 1 below.
• Ruff and git diff --check are clean.
• The new test_issue842 on main fails to import checkpoint_visible_nf4_linear.
• Six targeted mutations against the four CPU files:
  • the two correctness pins above are killed;
  • all four routing mutations survive on CPU. That is expected, because routing only exists on CUDA. What matters is which of them your GPU tests can see (item 4).

Please fix before the GPU run

1. The fragment turns CI red. changelog.d/0.75.1/1143.changed.md fails validate_fragments on two counts: it does not start with - and it does not mention #1143. So tests/test_issue487_changelog_fragments.py::test_repository_fragments_match_the_newest_release fails on this head and on the head merged with main, while main alone passes it (21 passed). It is also the one failure in each of the 9 CI test cells on this head; everything else in those cells passes. Please also credit yourself, for example:
   - **Streamed CUDA NF4 now follows bitsandbytes' own fused-versus-dequantise dispatch, so it runs the same kernels as resident NF4 (#842 by @Amix29 in #1143).** …
2. test_issue385 still uses the old reference. The ruling named this class. TestFloat16StreamingIsBitExact still calls install_dequant_forward(base) on its resident model (test_issue385_stream_dtype.py:499). Its docstring (:441-447) still says that comparing against the native path "would measure kernel-path differences rather than layer streaming", which the ruling makes false. Please drop the patch so the reference is plain bitsandbytes, and rewrite the docstring. With your change it should still read 0.0, because the Pin the bitsandbytes fused-vs-dequant NF4 divergence that #857 rests on #968 rewrite already asserts that patched == unpatched at the same fused M. The GPU run will confirm it.
3. "A test that fails when they move" is GPU-only for now. test_missing_bnb_private_dispatch_symbols_fail_toward_the_safe_function goes red if _gemm_4bit_use_custom_fn moves (the monkeypatch.delattr raises), but not if only _gemm_4bit_custom_max_m moves. Then only the training-M arm of the streamed-vs-resident test notices, and CI has no GPU, so it sees neither. from bitsandbytes.backends.cuda import ops imports here with CUDA hidden, and both names are present. A CPU test that runs importorskip on that module and asserts hasattr for both names makes the bitsandbytes upgrade PR itself turn CI red.
4. Two optional additions.
   • A CPU routing test (a fake CUDA x plus bnb's three conditions over a small grid) would pin routing in CI. One mutation is visible to no test at all, GPU included: dropping the max-M cap and the K % blocksize check (:367). The fixture never has M > 1536 or a misaligned K, and the only effect is the slowdown that ruling (a) removed.
   • In the streamed-vs-resident GPU test, expects_function=False at seq 128 encodes the 5070's window. My dispatch run agrees with it for sm_120, but deriving the expectation per call from bnb's decision would keep the test valid on other cards.

What the ruling means for this PR now

• Met on CPU evidence:

  • (a), routing from bnb's own dispatch, with the safe fallback when the private names are missing;
  • (c), the Pin the bitsandbytes fused-vs-dequant NF4 divergence that #857 rests on #968 rewrite and the four comments;
  • the ownership assertion in backward;
  • the docs paragraph at docs/performance-and-quantization.md:358-366, which names the CUDA path and the VRAM/time trade. The sentence at :434 is now accurate as written.

• Not met yet: the fragment (item 1), test_issue385's reference (item 2), and a CI-visible version of (a)'s "fails when they move" test (item 3).

• Only the GPU run can settle these: (b) itself, meaning the counters, the recycled-vs-private test and the plain-ctx arm, and streamed vs unpatched resident at the fused and training M. The same run settles the numerical half of (a) and (c), and whether the three CUDA mutations from my first review are killed. The files to run are:

  • the gpu subset of test_issue842;
  • test_issue968;
  • test_issue385 after item 2;
  • test_v07202::TestNF4ParityOnCuda;
  • test_v07300::TestBitsandbytesDispatchAtTrainingShape.

  The maintainer runs them on the RTX 5070 when the card is free.

Your description also lists #842's own measurements (throughput, peak VRAM, the large-layer alias control) as still open. Those belong to the issue's acceptance criteria.

The required checks on 8158f84c have run: lint and the smoke and floor jobs pass, and the 9 test cells fail on item 1 alone. The branch merges cleanly with main, #1141, #1178, #1196 and #1080.

[Amix29]

Addressed the 2026-09-25 review at head 8ea2f2d, rebased onto current main. The fragment now satisfies the release validator and credits #842/#1143; #385 compares streamed NF4 with native resident bitsandbytes; both private dispatch symbols are pinned on CPU-visible CI where the CUDA ops module imports; a CPU routing grid pins dispatch/max-M/alignment; and the GPU streamed-vs-resident test derives its expected route from bitsandbytes on the actual card. Local validation: 181 passed, 24 expected platform/GPU skips; Ruff and git diff --check clean. No CUDA result is claimed from this Mac.

[MakazhanAlpamys]

@Amix29 thank you. This round closes the three items from my review of 8158f84c. I checked each one by execution, not by reading the diff.

• The fragment. It now starts with - , credits (#842 by @Amix29 in #1143), and passes test_issue487. Dropping in #1143 from it turns that test red.
• test_issue385. It now compares against a plain bitsandbytes resident model, and the docstring says why.
• The private-name guard is visible to CI. I simulated a bitsandbytes release that moved each name, by deleting it from the imported bitsandbytes.backends.cuda.ops before the tests ran:
  • without _gemm_4bit_custom_max_m: test_bitsandbytes_private_dispatch_contract_is_visible_to_cpu_ci fails;
  • without _gemm_4bit_use_custom_fn: it fails too;
  • the control passes.
    In the CI logs of this head, that test runs and passes in the three ubuntu cells and skips on windows and macos, where the CUDA ops module does not import. A bitsandbytes upgrade PR therefore turns three cells red, which is what the item asked for.
• The CPU routing grid. It kills five routing mutations: the gate always off, the gate ignoring bitsandbytes, dropping the max-M cap, dropping the K % blocksize check, and swapping N and K in the heuristic call.

I ran everything below at 8ea2f2d0 on CPU (Windows, torch 2.14.0 with CUDA hidden, bitsandbytes 0.50.2).

Tests I ran

• test_issue842, test_v07300, test_v07202, test_issue968, test_issue385 and test_issue487: 183 passed, 22 skipped. That is the same 205 tests as your 181 / 24 on the M4 Max.
• The same files on your head merged with current main (20ab4668, which now includes fix(layer-stream): prefetch embed_tokens at the backward tail, not the next step's prime (#975) #1178's change to layer_stream_runtime.py): 183 passed, 22 skipped. The merge is clean. fix(layer-stream): prefetch embed_tokens at the backward tail, not the next step's prime (#975) #1178's own test_issue975… plus test_v07204 and test_v07200 on that tree: 277 passed, 18 skipped.
• test_issue842 on main: fails at collection, because checkpoint_visible_nf4_linear does not exist there.
• Ruff and git diff --check are clean. All 13 required checks are green on this head.

Mutations: 19 against those files, 12 killed

• The correctness pins hold:
  • holding the packed and quant tensors on ctx instead of save_for_backward (the Layer streaming: NF4 above ~165 MiB/layer produces silently wrong gradients (forward stays bit-exact) #331 class);
  • skipping the ownership check in backward;
  • returning no grad_x;
  • swapping absmax / absmax_8bit in the nested forward;
  • both Layer streaming: cut the ~10 ms fixed cost per layer visit, which dominates the step below ~128 tokens #841 cache mutations (never hit, or ignoring the mapping identity).
• Four survivors are CUDA-only by construction, so only the card can see them:
  • missing private names routed to dequantise instead of the Function;
  • the owner comparison inverted;
  • the fused arm sent back to bnb.matmul_4bit;
  • install_dequant_forward never consulting the gate.
    The committed gpu tests should kill all four on the card.

Three gaps worth closing on CPU before the card run, so the card only has to run the final head once. None blocks the draft.

1. The production ownership wiring is not pinned anywhere (layer_stream_runtime.py:1513). Deleting weight._soup_stream_owner_check = self._assert_nf4_slot_owner passes every CPU test and every committed gpu test:

   • test_backward_refuses_a_recycled_stream_slot attaches the check by hand;
   • the GPU streamed test never asserts that a check ran, and _nf4_stream's default of 2 layers over 2 buffers never recycles a slot within a step.

   On CPU, bitsandbytes' gemm_4bit is dequantise + linear, so the Function can be forced on in a recycling stream and compared with an unpatched resident model. This passes at your head and on the merged tree. It fails when the attachment is deleted, and when the owner comparison at :1522 is inverted:

   def test_forced_function_in_a_recycling_stream_matches_unpatched_resident(
       tmp_path, monkeypatch
   ):
       from test_v07202 import (
           _nf4_stream, _randomise_lora_b, _resident_nf4, _sync_adapters,
       )
   
       import soup_cli.utils.layer_stream_runtime as runtime_module
   
       calls = {"function": 0, "owner_checks": 0}
       real_function = runtime_module.checkpoint_visible_nf4_linear
   
       def counting_function(x, packed, quant_state, bias=None):
           calls["function"] += 1
           check = getattr(packed, "_soup_stream_owner_check", None)
           assert check is not None, "streamed NF4 weight carries no slot-owner check"
   
           def counted_check():
               calls["owner_checks"] += 1
               check()
   
           packed._soup_stream_owner_check = counted_check
           try:
               return real_function(x, packed, quant_state, bias)
           finally:
               packed._soup_stream_owner_check = check
   
       monkeypatch.setattr(runtime_module, "checkpoint_visible_nf4_linear", counting_function)
       monkeypatch.setattr(
           runtime_module, "_can_use_checkpoint_visible_nf4_gemm", lambda *_a, **_k: True
       )
   
       # 4 layers over 2 buffers: every slot is refilled by another layer mid-step.
       streamed, runtime, weights, _, _ = _nf4_stream(tmp_path, n_layers=4, buffers=2)
       resident = _resident_nf4(weights)
       _randomise_lora_b(resident)
       assert _sync_adapters(streamed, resident) > 0
   
       ids = torch.randint(0, 64, (1, 16), generator=torch.Generator().manual_seed(5))
       batch = {"input_ids": ids, "attention_mask": torch.ones_like(ids), "labels": ids}
       streamed.train()
       resident.train()
       streamed_loss = streamed(**batch).loss
       streamed_loss.backward()
       resident_loss = resident(**batch).loss
       resident_loss.backward()
   
       assert calls["function"] > 0 and calls["owner_checks"] > 0
       assert runtime.pool.loads > 4, "no slot was recycled"
       assert torch.equal(streamed_loss, resident_loss)
   
       def lora_grads(model):
           return {
               name.replace(".inner.", "."): param.grad
               for name, param in model.named_parameters()
               if "lora_" in name and param.grad is not None
           }
   
       mine, theirs = lora_grads(streamed), lora_grads(resident)
       assert mine and mine.keys() == theirs.keys()
       for name, grad in mine.items():
           assert torch.equal(grad, theirs[name]), name

   I ran it as a standalone file that puts tests/ on sys.path; inside test_issue842_checkpoint_visible_nf4.py the existing sys.path line already covers the import. For the same reason, please consider n_layers=4 in test_streamed_matches_unpatched_resident_logits_and_all_lora_grads. Then the fused-M arm also runs with real slot recycling on CUDA, where Layer streaming: NF4 above ~165 MiB/layer produces silently wrong gradients (forward stays bit-exact) #331's asynchronous-copy race actually lives.

2. M from 3-D activations (:364). Taking m from x.shape[0] instead of x.numel() // k passes everything, CPU and GPU. Every grid row is 2-D, and the GPU test derives its expectation from Soup's own gate, so it agrees with itself. A real (1, 128, H) activation would read as M=1 and take the Function at training M. That is the third dequantisation ruling (a) removed, and since the numerics do not change, no equality test notices. Two 3-D rows in your grid close it. They pass at your head and fail under that mutation:

   @pytest.mark.parametrize(
       ("shape", "expected", "expected_m"),
       [((1, 97, 64), False, None), ((2, 3, 64), True, 6)],
       ids=["3d-above-cap", "3d-inside-window"],
   )
   def test_m_counts_every_leading_dimension(monkeypatch, shape, expected, expected_m):
       from types import SimpleNamespace
   
       import bitsandbytes.backends.cuda as bnb_cuda_package
   
       from soup_cli.utils.layer_stream_runtime import _can_use_checkpoint_visible_nf4_gemm
   
       calls = []
   
       def use_custom_fn(device_index, dtype, rows, out_features, in_features):
           calls.append((device_index, dtype, rows, out_features, in_features))
           return True
   
       fake_ops = SimpleNamespace(
           _gemm_4bit_custom_max_m=96, _gemm_4bit_use_custom_fn=use_custom_fn
       )
       monkeypatch.setitem(sys.modules, "bitsandbytes.backends.cuda.ops", fake_ops)
       monkeypatch.setattr(bnb_cuda_package, "ops", fake_ops, raising=False)
       numel = shape[0] * shape[1] * shape[2]
       x = SimpleNamespace(
           device=SimpleNamespace(type="cuda", index=0), dtype=torch.bfloat16,
           shape=shape, numel=lambda: numel,
       )
       state = SimpleNamespace(nested=False, shape=(128, 64), blocksize=64)
       assert _can_use_checkpoint_visible_nf4_gemm(x, state) is expected
       assert calls == ([] if expected_m is None else [(0, torch.bfloat16, expected_m, 128, 64)])

3. The safe direction for missing names is pinned only by a gpu test (:357-361). Returning False there passes on CPU. The same fake-module trick with SimpleNamespace() (neither name present), a CUDA-like x of shape (512, 64) and an expected True makes it CI-visible. It fails under that mutation.

A cosmetic fourth: the nested state2.blocksize != 256 guard (:345) can be removed without a failure. It is defensive, since bitsandbytes' own MatMul4Bit raises for any other nested block size and its quantiser always uses 256. I would leave it as it is.

What remains before this leaves draft

• The CUDA run of ruling (b), on the final head merged with main:
  • Command: pytest -m gpu --no-cov -v tests/test_issue842_checkpoint_visible_nf4.py tests/test_issue968_bnb_nf4_path_divergence.py "tests/test_issue385_stream_dtype.py::TestFloat16StreamingIsBitExact" "tests/test_v07202.py::TestNF4ParityOnCuda" "tests/test_v07300.py::TestBitsandbytesDispatchAtTrainingShape".
  • That selects 19 tests. The bar is 19 passed. It includes test_issue385's [nf4] cases reading exactly 0.0 in both fp16 and bf16, which is the old Streamed NF4 is not bit-exact vs resident NF4 on Blackwell (sm_120): 4.9e-04 / 3.9e-03, while quantization=none stays exact #776 failure.
  • Then the three mutations from my first review, one at a time, each of which must turn at least one of those tests red:
    • the gate off;
    • the fused arm sent back to bnb.matmul_4bit;
    • the recompute through dequantise + F.linear.
  • This needs the RTX 5070, not your Mac.
• The PR body. Its "Still required before leaving draft" section still lists Layer streaming: a checkpoint-visible fused NF4 dequantise-and-multiply (upper bound 9.8% of the step) #842's own measurements: throughput, peak VRAM, and the aliasing control above ~165 MiB per layer. The note of 09-25 placed those on the issue's acceptance criteria rather than on this PR. If you agree, that section can say so. The PR already says "Refs Layer streaming: a checkpoint-visible fused NF4 dequantise-and-multiply (upper bound 9.8% of the step) #842", not "Closes".

[Amix29]

Closed the three CPU gaps from the latest review at 433b35a4, rebased on current main:

• production owner-check wiring is now exercised through a real 4-layer / 2-buffer recycling stream; the test proves Function and owner-check execution, buffer recycling, exact resident loss, and every LoRA gradient;
• 3-D activations pin M = numel / K both above and inside the custom-kernel window;
• missing private dispatch names pin the safe Function direction in CPU-visible CI;
• the CUDA streamed-vs-resident case now also uses 4 layers / 2 buffers.

Local M4 Max validation: 185 passed, 24 expected platform/GPU skips across the six focused files; the #842 module alone is 15 passed, 5 skipped. Ruff, preview E30, and git diff --check pass. I also updated the PR body so its remaining PR-level gate is only the named CUDA acceptance run; #842 performance/VRAM/large-layer measurements remain issue-level work. No CUDA result is claimed from this Mac.

[MakazhanAlpamys]

@Amix29 thank you. 433b35a4 closes all three CPU gaps from my review of 8ea2f2d0, and I checked each one by breaking the code and watching your new test go red, not by reading the diff.

• The production ownership wiring. test_forced_function_in_a_recycling_stream_matches_unpatched_resident fails when the _soup_stream_owner_check attachment is deleted (layer_stream_runtime.py:1555) and when the owner comparison is inverted (:1564). It also does more than I asked for. Sending the fused arm back to bnb.matmul_4bit (:438) and never consulting the gate in install_dequant_forward (:437) now fail on CPU too, because the Function counter reads 0. At 8ea2f2d0 only the card could see those two, and now every CI cell can.
• M from 3-D activations. test_cuda_route_flattens_all_leading_dimensions_on_cpu fails when m is taken from x.shape[0] (:364), in both of its rows.
• The safe direction for missing names. test_missing_bnb_private_dispatch_symbols_fail_toward_safe_function_on_cpu fails when the missing-name branch returns False (:359-361).
• The GPU streamed-vs-resident test now uses 4 layers over 2 buffers, and the PR body now names only the CUDA acceptance run as the PR-level gate. Both are what I asked for.

I also checked the rebase. git range-diff shows your earlier commits unchanged except for the docs context line, where main's #1178 rewrote the "How it works" paragraph. The new NF4 paragraph now follows it and reads correctly.

Everything below ran at 433b35a4 on CPU (Windows, torch 2.14.0 with CUDA hidden, bitsandbytes 0.50.2).

Tests I ran

• test_issue842, test_v07300, test_v07202, test_issue968, test_issue385 and test_issue487: 187 passed, 22 skipped. That is the same 209 tests as your 185 / 24 on the M4 Max.
• The same files on your head merged with current main (df106781): 187 passed, 22 skipped. The merge is clean.
• Together with feat(stream): prepare real preference VRAM probes (#840) #1141, which also touches layer_stream_runtime.py: both PRs' review files plus test_issue975… on the merged tree give 327 passed, 42 skipped.
• test_issue842 on main fails at collection, because checkpoint_visible_nf4_linear does not exist there.
• Ruff and git diff --check are clean, and the fragment still passes the changelog validator.

Mutations: 20 against those files, 18 killed

• Every earlier pin still holds: the routing grid (gate off, gate ignoring bitsandbytes, no max-M cap, no K % blocksize check, N/K swapped), save_for_backward vs ctx (the Layer streaming: NF4 above ~165 MiB/layer produces silently wrong gradients (forward stays bit-exact) #331 class), no grad_x, the swapped nested absmax, both Layer streaming: cut the ~10 ms fixed cost per layer visit, which dominates the step below ~128 tokens #841 cache mutations, and the fragment.
• The six that could be fixed without a card are now killed: missing names routed to dequantise, M from the leading dimension, the ownership attachment deleted, the owner comparison inverted, the fused arm sent to bnb.matmul_4bit, and the gate never consulted.
• Two survive. One is the nested state2.blocksize guard (:347), which I said to leave as it is. The other is new, and it is optional.

One optional addition: the check's refusing branch. Your test proves that the check is attached and stays silent for the right owner. Nothing yet proves that _assert_nf4_slot_owner (:1561) raises when the slot holds another layer. test_backward_refuses_a_recycled_stream_slot uses a hand-written check, and no GPU test recycles a slot mid-backward. So replacing if owner != self.idx: with if False: passes all 209 CPU tests. This passes at your head and fails under that change ("DID NOT RAISE"). The module's existing sys.path line covers the import:

def test_the_production_owner_check_refuses_a_recycled_slot(tmp_path):
    from test_v07202 import _nf4_stream

    import soup_cli.utils.layer_stream_runtime as runtime_module

    model, runtime, _, _, _ = _nf4_stream(tmp_path, n_layers=4, buffers=2)
    layer = runtime_module.decoder_owner(model).layers[0]
    slot = runtime.pool.slot_for(layer.idx)

    runtime.pool.owner[slot] = layer.idx
    layer._assert_nf4_slot_owner()  # the right owner: silent

    runtime.pool.owner[slot] = layer.idx + 2  # layer 2 has refilled the slot
    with pytest.raises(RuntimeError, match="recycled weight slot"):
        layer._assert_nf4_slot_owner()

CI. The required checks on 433b35a4 have been queued since 11:17Z and have not started yet, so there is no CI result for this head. Your new CPU tests use the same imports and the same CPU gemm_4bit op as tests that passed in all nine cells on 8ea2f2d0, but I have not seen them run on macOS or Windows.

What remains before this leaves draft. Nothing more is needed from you unless you want to add the test above. The rest needs the RTX 5070, not your Mac:

• the 19 gpu tests, on the final head merged with main: pytest -m gpu --no-cov -v tests/test_issue842_checkpoint_visible_nf4.py tests/test_issue968_bnb_nf4_path_divergence.py "tests/test_issue385_stream_dtype.py::TestFloat16StreamingIsBitExact" "tests/test_v07202.py::TestNF4ParityOnCuda" "tests/test_v07300.py::TestBitsandbytesDispatchAtTrainingShape". The bar is 19 passed, including test_issue385's [nf4] cases reading exactly 0.0 in fp16 and bf16;
• the three mutations from my first review, one at a time, each turning at least one of those tests red: the gate off, the fused arm sent back to bnb.matmul_4bit, and the recompute through dequantise + F.linear.

If you do add the test, please add it before that run, so the card only runs the final head once.

[Amix29]

Added the optional production refusal pin before the CUDA run at e7e6476d.

The new test builds the real 4-layer / 2-buffer NF4 stream, calls the production _assert_nf4_slot_owner, proves it is silent for the current owner, then simulates layer 2 recycling layer 0s slot and requires the documented RuntimeError. It fails if the production refusing branch is disabled.

Local M4 Max validation:

• Layer streaming: a checkpoint-visible fused NF4 dequantise-and-multiply (upper bound 9.8% of the step) #842 module: 16 passed, 5 expected GPU/platform skips;
• six focused compatibility files: 186 passed, 24 expected skips;
• Ruff, preview E30, and git diff --check pass.

The PR body is refreshed. The draft now has no remaining CPU-side review item; the RTX 5070 acceptance run remains external, and no CUDA result is claimed here.

[MakazhanAlpamys]

@Amix29 thank you. e7e6476d adds the refusal pin, and it arrived before the card run, as I asked, so the RTX 5070 only has to run this head once.

I checked it by breaking the code, not by reading the test:

• The survivor is closed. With if owner != self.idx: replaced by if False: (layer_stream_runtime.py:1564), your new test_the_production_owner_check_refuses_a_recycled_slot fails with "DID NOT RAISE RuntimeError": 187 passed, 1 failed. At your head it passes.
• Three more mutations of the same check, all caught:
  • With the comparison inverted, your new test and the forced-Function recycling test both fail (186 passed, 2 failed).
  • When the refusal says "stale weight slot" instead of "recycled weight slot", your test's match fails (1 failed).
  • When the check reads slot 0 for every layer, the forced-Function recycling test fails (1 failed). Your new test cannot see that one, because layer 0's slot is 0, but it is already covered.

Every mutation was undone, and git diff was empty afterwards. src/ has not changed since 433b35a4, so the 18 mutations killed there still stand. The only survivor on these files is the nested state2.blocksize guard, which stays as it is.

Tests I ran (Windows, CPU with CUDA hidden, bitsandbytes 0.50.2):

• test_issue842, test_v07300, test_v07202, test_issue968, test_issue385 and test_issue487 at e7e6476d: 188 passed, 22 skipped. That is the same 210 tests as your 186 passed / 24 skipped on the M4 Max.
• The same files on your head merged with current main (58a57521): 188 passed, 22 skipped. The merge is clean.
• Ruff and git diff --check are clean.

CI. The required checks on e7e6476d have been queued since 12:29Z and have not started. The run on 433b35a4 was cancelled by this push, so the tests added in the last two commits have not run on CI yet. The last completed run is still 8ea2f2d0, which was green.

What remains. Nothing CPU-side. This PR now waits only for the CUDA acceptance run on the RTX 5070, on this head merged with main: the -m gpu command from my earlier review, plus the three mutations (gate off, fused arm sent back to bnb.matmul_4bit, and the recompute through dequantise + F.linear), each of which must turn at least one of those tests red.

[MakazhanAlpamys]

@Amix29, both GPU gates are done, on e7e6476d merged with main (RTX 5070, sm_120):

• G1, the 19-item -m gpu command: 19 passed, including test_issue385 [nf4] in fp16 and bf16 at exactly 0.0, so the Streamed NF4 is not bit-exact vs resident NF4 on Blackwell (sm_120): 4.9e-04 / 3.9e-03, while quantization=none stays exact #776 divergence does not reproduce on this head.
• G2, one mutation at a time, each restored byte for byte afterwards:
  • gate off: 7 failed;
  • recompute through dequantize_4bit + F.linear: 1 failed. Only 842's streamed-vs-unpatched [fused-m] test sees it: the logits still match exactly, but layer 0's q_proj.lora_A grad differs (-1.3513e-06 vs -1.3619e-06);
  • fused arm sent back to bnb.matmul_4bit: 3 failed, caught by the matmul_4bit and Function counters.

Every mutation turns at least one G1 test red, which was the bar. One thing worth knowing: exactly one of the 19 tests sees the recompute mutation.

From my side this is ready to merge. Mark it ready for review when you want it in, and it merges once the required checks are green on that head. #1258 still overlaps this PR in layer_stream_runtime.py, the #975 test and the performance docs, so whichever lands second re-resolves.

[MakazhanAlpamys]

@Amix29, thanks for marking this ready. #1258 merged first (15a66781), and it touches the same three files as this PR: src/soup_cli/utils/layer_stream_runtime.py, tests/test_issue975_backward_tail_embed_prefetch.py and docs/performance-and-quantization.md. So this one re-resolves, as we said. Please rebase on main. Since the rebase changes layer_stream_runtime.py, I'll re-run the G1 set and the three G2 mutations on the card on the new head before merging, then merge on green.

[Amix29]

Rebased on current main after #1258 as requested. New head: e1bcb669. The six commits replayed cleanly, including the shared layer_stream_runtime.py, tail-prefetch test context, and performance documentation.

Focused validation across the #842/NF4 compatibility files: 166 passed, 23 expected GPU/platform skips. Full Ruff and git diff --check are clean. The maintainer-side RTX 5070 G1 and three G2 mutation reruns remain the hardware gate; no CUDA result is claimed from this M4 Max.

[MakazhanAlpamys]

@Amix29 thank you for the quick rebase. I checked it before running anything on the card:

• git range-diff against the pre-rebase series shows five of your six commits unchanged. The sixth differs only in the docs context line, because #1258 rewrote the "How it works" paragraph that your NF4 paragraph follows. Your paragraph still reads correctly after it.
• The PR's diff over src/, tests/ and changelog.d/ has the same stable patch-id before and after the rebase. So the code, the tests and the fragment are byte-for-byte what ran on the card on 09-27. #1258's change to layer_stream_runtime.py (the head-prefetch timing) does not overlap your hunks.

On the RTX 5070 (sm_120, torch 2.14.0+cu130, bitsandbytes 0.50.2), on e1bcb669 merged with current main (dccc8680). The merge is clean.

• G1, the 19-item -m gpu command: 19 passed, none skipped. That includes test_issue385's [nf4] cases at exactly 0.0 in fp16 and bf16, and your 4-layer / 2-buffer streamed-vs-unpatched-resident test at the fused and the training M.
• G2, one mutation at a time, each restored byte for byte afterwards:
  • gate off: 7 failed;
  • recompute through dequantize_4bit + F.linear: 1 failed, again only the [fused-m] streamed-vs-resident test. The logits still match, and layer 0's q_proj.lora_A grad differs, -1.3513e-06 vs -1.3619e-06, the same digits as before the rebase;
  • fused arm sent back to bnb.matmul_4bit: 3 failed, caught by the matmul_4bit counter (52 calls where 0 are allowed) and the Function counter.
• Without the fix, meaning layer_stream_runtime.py taken from main with your tests kept, test_issue385 [nf4] reads 4.8828125e-04 in fp16 and 3.90625e-03 in bf16. Those are the #776 numbers exactly, so on this card your change is what brings them to 0.0.

On CPU (Windows, CUDA hidden), on the same merged tree:

• The six review files (test_issue842, test_v07300, test_v07202, test_issue968, test_issue385, test_issue487): 188 passed, 22 skipped, the same 210 tests as before the rebase.
• The streaming files around it: #1258's two test_issue1174_* files, test_issue975…, both test_issue379_* files (they drive install_dequant_forward), test_issue361_nf4_throughput and test_v07200. 352 passed, 6 skipped.
• Seven mutations on the rebased code, all killed: the production owner check that never refuses, the owner-check attachment deleted, save_for_backward replaced by plain ctx attributes, the view cache never hit, the gate off, M taken from x.shape[0], and the fragment without in #1143. Each was undone, and git status was empty afterwards.
• Ruff and git diff --check are clean. Your new test file fails at collection on main, as it should.

This is ready to merge. It now waits only on the 13 required checks, which have been queued since 14:53Z and have not started yet. It closes no issue: #841 is already closed as folded into #842, and #842's own measurement criteria stay open there.

[MakazhanAlpamys]

Merged as 33639fb4. Thank you, @Amix29: you made Soup ask bitsandbytes for its routing decision instead of copying one card's window into the code, and you had every gap review found pinned on CPU before the card ran, so after the rebase it only had to confirm. On the RTX 5070, on the head merged with main: 19 of 19 gpu tests passed, including [nf4] at exactly 0.0 in fp16 and bf16, and each of the three ruling mutations turned tests red (7, 1 and 3). With the source taken from main, the #776 numbers came back. On CPU: 188 passed / 22 skipped over the six review files, 352 / 6 over the neighbouring streaming files, and 7 of 7 mutations killed.