[mcencini]

What changed

execution="auto" is the default for simulators, estimators and ModelOperator. A call too small to repay a launch, or one carrying an autograd graph, runs where its tensors are. A larger call is cut into voxel chunks sized to each card's free memory less a 512 MiB reserve and dealt to every visible GPU. This works whether the volume starts on the host or on a card, and the answer is written back on the input's device.

Every pass (forward, jvp, vjp, vjp_jvp) and every physics model is chunked: shaped-pulse slice profiles, per-voxel rotation pairs (pTx dynamic transmit), MT lineshape, exchange and tabulated pools, shims and diffusion. A backward pass runs under the policy its forward pass ran under. ModelOperator sizes its chunks on the full time series before projection.

Why

Matteo asked for this in the thread "blochsim host/device chunking audit": keep the trimmed public API from #41, but run large problems with a bounded VRAM footprint over both cards, the way bartorch's Toeplitz NUFFT does. Before this change:

• device-resident volumes were never chunked;
• two cards helped only when there were several trains;
• profiles, pools, MT, exchange and pTx were never streamed.

The audit and the agreed plan are in the project plan blochsim-execution-audit.md.

Validation

• pytest tests -n auto: everything passes on CPU. The one exception is a 3-second timing test in tests/optim/test_design.py, which failed under xdist contention and passes on its own.
• New tests/sequence/test_chunks.py compares the whole volume with the chunked volume on all four passes. It covers echo-train, spoiled and shimmed volumes, tabulated pools (2f, 3s/profile, 4s/dynamic, 4s/shimmed), the slice profile, the dynamic pair, and the GPU kernels through _gpu_host. Chunk widths are 1, 2, 4, 7 and 64, over 1 to 3 devices.
• The FSE design loop from test_design.py takes 60 Adam iterations at 0.59 to 0.62 s on this branch against 0.50 to 0.61 s on main, best of 3, which is within noise. Small design calls with a graph never reach the chunker.
• Not run: CUDA. The container has no GPU, so the CUDA legs skip. Please check on the 2× RTX 4000 Ada:
  • pytest tests/sequence/test_chunks.py tests/sequence/test_offload.py tests/sequence/test_execution.py;
  • peak VRAM on a large resident volume;
  • the virtual scanner;
  • the pypulseqpp FSE design.

Documentation impact

• docs/explanations/execution.md: the TL;DR, Capabilities, Division of responsibility and Method sections are rewritten for the new default.
• The routes figure now shows two routes.
• AGENTS.md and docs/developer-guide/documentation.md no longer name the private knobs.
• There is a CHANGELOG entry.

Checklist

• pytest tests/ passes locally, and a new test covers what changed.
• pre-commit run --all-files is clean: ruff format and ruff check both pass.
• Public functions and classes carry a numpydoc docstring.
• Comments and docstrings describe the code as it is now.
• A change to what the kernels compute is in both of them. The kernels themselves are unchanged; chunk dispatch is Python around them.
• The documentation is updated where the change is visible to a caller, and bash scripts/build_docs.sh still builds. The build is running now.

🤖 Generated with Claude Code

https://claude.ai/code/session_01PrEk8PfzyAojw2pKzoTiHd

---

Generated by Claude Code