Merge quantem-cuda kernel dispatch (upstream PR #243) onto the benchmark stack#21
Open
cedriclim1 wants to merge 24 commits into
Open
Merge quantem-cuda kernel dispatch (upstream PR #243) onto the benchmark stack#21cedriclim1 wants to merge 24 commits into
cedriclim1 wants to merge 24 commits into
Conversation
Add a has_quantem_cuda capability flag (mirroring has_torch/has_cupy) and a use_cuda_kernels config option (default true) as a kill-switch. ObjectPixelated.get_tv_loss now goes through tv_loss_vol_sq in tomography/utils.py, which dispatches to quantem-cuda's fused squared-anisotropic TV kernel (identical math, analytic backward, 10-14x faster fwd+bwd at 256^3-512^3) when the package is installed, the object is a CUDA fp32 tensor, and the kill-switch is on; otherwise it falls back to the existing pure-torch expression. quantem-cuda is not yet on PyPI, so the [cuda] extra is deferred to a one-line follow-up after the 0.1.0 publish (an unpublished package in optional-dependencies breaks uv lock for everyone). Until then, installing quantem-cuda from source enables the dispatch. Tests cover the fallback paths (CPU, no quantem-cuda, kill-switch) and kernel parity for values and gradients when it is installed.
interpolate_ms_features_tilted routes each multiscale level through the fused kernel (rotate + grid_sample + Hadamard product, analytic backward) when quantem-cuda is installed, tensors are CUDA fp32, and the use_cuda_kernels config flag is on; the torch path is unchanged and remains the fallback. Covered by GPU/CPU parity, gradient parity, per-scale dispatch-spy, and kill-switch tests.
Evaluate the four TV taps (base, +x, +y, +z) in one batched 4N-point forward instead of four serial 10k-point calls, which were kernel-launch bound (volume TV 1.62 -> 0.52 ms, full training step 7.16 -> 5.23 ms at the 200^3 benchmark config). Also gate the plane and volume TV terms on their own weights: tv_plane > 0 with tv_vol == 0 previously applied no plane TV at all, and tv_vol > 0 computed a zero-weighted plane TV. The batched call uses the unwrapped model, consistent with the base tap and _get_plane_tv_loss.
Each training step previously made two model calls on a tensor-decomp object: the main batch forward and a second 40k-point call for the volume-TV finite-difference taps. Both backwards accumulate gradients into the same plane parameters, so autograd runs the full gradient-accumulation traffic twice; profiling showed that traffic (copy_/add_/fill_, ~1.8 ms/step) exceeds the interpolation backward itself. The object model now exposes sample_tv_tap_coords(), the training loop concatenates the returned tap points onto the main batch for a single model call via forward_with_tv_taps() (mask and hard constraints apply to the main chunk only; tap densities stay raw, matching the previous TV semantics), and the tap densities reach get_volume_tv_loss through ReconstructionContext. The two-call path remains as a fallback for direct callers and non-tensor-decomp models. Full step 5.67 -> 4.79 ms (-15.5%) at the 200^3 benchmark config; loss and gradient parity verified against the two-call path, including taps crossing the [-1, 1] boundary.
Fold volume-TV tap evaluation into the main forward pass
quantem-cuda moved its kernels into per-module submodules mirroring quantem: the K-Planes interpolation now lives in quantem.cuda.core.ml (KPlanesTILTED is a quantem.core.ml model shared across applications) and the TV kernels in quantem.cuda.core. Update the dispatch imports and the test monkeypatch targets accordingly.
docs/notebooks/quantem_cuda_kernels.ipynb spells out exactly what the quantem-cuda TV kernels implement (tv_loss_sq_3d: unnormalized squared anisotropic sum, tv_vol parity; tv_loss_iso_3d: corner-restricted isotropic mean, eps inside the sqrt) and what they do not (no anisotropic-L1 kernel, no per-axis weights, iso needs D >= 2), with parity checks, a ptychography section (complex multislice objects, per-axis weighting recipe, honest op-level timings), the fused TILTED K-Planes interpolation contract, and the transparent-dispatch config surface. Executed on an RTX PRO 6000 so the outputs are visible in review. .gitignore gains an exemption for curated docs notebooks.
…kernel ObjectConstraints._calc_tv_loss now routes fp32 CUDA 3-D arrays through quantem.cuda.core.tv_loss_l1_3d (per-axis |diff| sums), composing the same anisotropic-L1 functional — per-axis (tv_weight_z, tv_weight_xy) weights and active-axis normalization included — with identical gradients (sign(0) = 0). Weighted size-1 axes fall through to torch so degenerate inputs behave exactly as before; the kill switch and capability flag are the same as the other dispatch points. ~4x fwd+bwd at multislice (16, 1024, 1024) sizes, neutral at single-slice sizes.
…ity gate docs/notebooks/quantem_cuda_kernels.ipynb now documents tv_loss_l1_3d (the ptychography functional) with a verified parity composition and updated timings/dispatch tables. docs/notebooks/ptycho_tv_kernel_quality.ipynb reruns the tutorial ducky reconstruction (simulated dataset, single-slice xy TV and multislice z+xy TV) with and without the kernel and gates on SSIM of the reconstructed phase against an fp-perturbation chaos floor: the pipeline is bit-exact deterministic, so a 1-ulp TV rounding difference diverges trajectories exactly like a 2e-7 relative defocus perturbation does. Both configs pass (SSIM 0.997 vs floor 0.996 single-slice; 0.992 vs 0.988 multislice), with final data-fidelity losses agreeing within the floor's scatter.
The repo gitignores *.ipynb on purpose — drop the docs/notebooks exemption and the two committed notebooks. They remain available as a PR attachment for review.
…p index tensors in interpolate_ms_features - interpolate_ms_features_cp_tilted rebuilt the (3T, 1, B, 2) sampling grid (reshape/permute plus a zeros_like and stack) inside the per-scale loop even though it only depends on the rotated points; build it once before the loop. - interpolate_ms_features (non-tilted) projected points onto the three planes with list-based advanced indexing, which allocates an index tensor on device (a host-to-device copy) three times per forward; use unbind/stack views, matching the TILTED variant. Both bitwise-identical to the previous implementations; ~1.3x each on the isolated forward (B=4096, 3 scales; CP with T=8).
- TomographyINRDataset now implements __getitems__: torch's DataLoader calls it with the whole batch of indices, replacing one Python __getitem__ call (plus per-item dict construction and collate stacking) per sample with a few tensor ops. setup_dataloader pairs datasets that define __getitems__ with a passthrough collate_fn (module-level so spawn workers can pickle it); the random_split Subset path delegates the batched form automatically. Batches are identical to default collate (values and dtypes). 39.7x faster batch fetching (4.16 -> 0.105 ms/batch, batch_size=1024, 60x256x256 stack). - The per-epoch loss reduction issued three all_reduce calls and three .item() host syncs; stack the three scalars into one tensor for a single all_reduce and a single .tolist() sync.
…ction - _build_optimizer now passes fused=True to Adam/AdamW when every parameter is on a CUDA device: the whole optimizer step runs in one fused kernel, 2.0x faster than the default foreach path on KPlanesTILTED-sized grids (1.24 -> 0.61 ms/step; ~9% end-to-end per training iteration, where opt.step was 19% of the profile). Same update rule; parameters agree to ~2e-6 after 5 steps (kernel-order float effects). - The epoch losses were re-wrapped in a tensor, all_reduced a second time (idempotent on already rank-averaged values), and synced to host again right after the stacked reduction; the redundant block is removed.
KPlanes(use_hybrid_mlp=False) -- the constructor default -- never built sigma_net, yet forward, get_params, and param_keys all reference it unconditionally, so a default-constructed KPlanes crashed with AttributeError on first use (including ObjectTensorDecomp.from_model in the reconstruction pipeline). KPlanesTILTED._build_sigma_net and CPTilted already fall back to a single linear head in this case; do the same here with identical initialization. Regression tests cover get_params/param_keys consistency and a finite forward pass on the default constructor.
reconnect_optimizer_to_parameters restored per-group hyperparameters by zip-index against a fresh get_optimization_parameters() call. Group order and membership from that dict are not contractual (the dataset model's groups change with alignment settings), so any difference between creation time and reconnect time silently attached the wrong learning rate to every group -- no error, just a reconstruction quietly training with swapped lrs. set_optimizer now records each group's key in the torch param group under "name" (torch preserves unknown keys), and reconnect re-joins hyperparameters by that key; groups with no old counterpart keep optimizer defaults. Optimizers restored from checkpoints predating group names fall back to the previous index alignment. Regression tests cover reorder, group addition, state survival, and the legacy fallback.
HSiren INR training is MLP-bound: a 256-wide sine MLP evaluated on ~1.6M points per batch, where eager mode pays full memory round-trips between every linear and sine. ObjectINR.from_model(..., compile_model=True) now routes the main forward through torch.compile. Design: the compiled artifact is cached in a module-level WeakKeyDictionary keyed by the live model object, so nothing extra lands on the module tree or in AutoSerialize output, and reset()/rebuild_model() (which swap the model object) invalidate it naturally. The TV autograd.grad recompute stays on the eager model (double backward through compiled graphs is not reliable). Default remains eager. 200-cube phantom, 29 tilts, batch 8192, RTX PRO 6000: 13.798/13.819 s/epoch eager -> 7.804/7.802 compiled, ~1.77x, matching losses. Isolated fwd+bwd microbenchmark: 1.73x.
…o bench/cuda-dispatch # Conflicts: # src/quantem/core/ml/models/kplanes.py # src/quantem/tomography/object_models.py # src/quantem/tomography/tomography.py # src/quantem/tomography/tomography_context.py
The PR-side helper was written against the pre-fix base and masked only x/y; the main batch routes through it whenever tv_vol > 0, which would silently revert the z-axis out-of-volume fix relative to feat/cv-pixel-holdout. Mirror forward()'s mask exactly.
Pre-existing break on feat/cv-pixel-holdout (not introduced by the dispatch merge): the box_fixed_ds commit turned integrate_rays into an instance method dispatching on ray_sampling, but the upstream test still called the old staticmethod form. First CI run on this lineage surfaced it. Pin ray_sampling='legacy' so the test keeps asserting the original step-size summation through the public API.
…to bench/cuda-dispatch
…' into bench/cuda-dispatch
…ame' into bench/cuda-dispatch
…ch reductions Keep the optimizer_mixin fused Adam/AdamW selection (all-CUDA param groups) and tomography.py's single stacked all_reduce + removal of the redundant second reduction. Drop the ddp.py passthrough collate and dataset __getitems__ vectorization: superseded by DeviceBatchSampler, which already builds GPU-resident batches — those hunks would be dead code on this stack.
…n' into bench/cuda-dispatch
This was referenced Jul 9, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Integrates the optional quantem-cuda kernel dispatch (electronmicroscopy#243,
feat/quantem-cuda-extra) onto the benchmark overlay stack (feat/cv-pixel-holdout) for the A1.1 speed-up study.What's here
68173f6— the 10-commit dispatch branch. Four files hit textual conflicts; the two semantic resolutions:interpolate_ms_features_tilted: the fused-kernel early return now appliesscale_gatesper scale (the PR path would otherwise silently bypass the coarse-to-fine prior).get_volume_tv_loss: composes the PR'sprecomputed_tap_densitiesfast path with ourtv_vol_anisoweighting;n = shape[0]//4works for both paths.d378f70— the PR-sideforward_with_tv_tapswas written against the pre-a16ea38 base and masked only x/y; restored the 3-axis out-of-volume mask (main batch routes through this helper whenevertv_vol > 0). Verified on A100: main chunk bit-exact vsforward()(max diff 0.0) including out-of-volume z.Validation (A100 / sm_80, Perlmutter)
tv_loss_sq_3d5.4× @ 200³.Full tables:
results/a11_paired/{summary.json,PR243_COMMENT_DRAFT.md}in the benchmark workspace.Note:
kplanes_tilted_tv_fuseexists in quantem-cuda main but is not wired up here (matches upstream electronmicroscopy#243); flagged as a follow-up iftv_volworkloads matter.Bundle: open fork PRs (added 2026-07-09)
Per-PR verdicts (reviewed against this branch's HEAD; #15 excluded by request):
b74fe46)3e82b26)__getitems__/collate hunks — superseded by DeviceBatchSamplerkplanes_fuse, which quantem-cuda does not export — guard makes it a permanent no-op; also bypasses theuse_cuda_kernelspolicya6159c2)use_cuda_kernelsand appliesscale_gates)e661f56)20f767b)forward_with_tv_tapsstays eager)Validation (A100): 266 passed / 12 skipped across ml + tomography + diffractive suites; 3-epoch smoke with kernel auto-enable OK; 30-epoch seed-0 run: 2.15 s/epoch, VRAM 2.66 GB.
Numerics note: the 30-epoch seed-0
mae_fullmoved from 0.008945 (pre-bundle) to 0.008055 — beyond seed noise, most plausibly the #19 reconnect fix correcting hyperparameter assignment (an improvement, but it means pre- and post-bundle runs should not be mixed in one comparison).