Skip to content

Dispatch interpolate_ms_features to the fused quantem-cuda kernel - #16

Closed
cedriclim1 wants to merge 1 commit into
feat/tomography-inr-fixesfrom
perf/kplanes-fuse-dispatch
Closed

Dispatch interpolate_ms_features to the fused quantem-cuda kernel#16
cedriclim1 wants to merge 1 commit into
feat/tomography-inr-fixesfrom
perf/kplanes-fuse-dispatch

Conversation

@cedriclim1

Copy link
Copy Markdown
Owner

What

interpolate_ms_features (non-tilted KPlanes) now dispatches to quantem.cuda.core.ml.kplanes_fuse — the fused fwd+bwd CUDA kernel added in cedriclim1/quantem-cuda#1 — when the inputs qualify (CUDA, fp32, standard [B,3] points and [3,C,H,W] grids). All other cases (CPU, fp64, missing quantem-cuda) fall through to the unchanged torch path; the import is guarded so installations without quantem-cuda see zero difference.

Why

F.grid_sample backward is the dominant cost of non-tilted KPlanes INR training. The fused kernel benchmarks at ~9.6x fwd+bwd on the isolated op; end-to-end this translates to:

200³ phantom, 29 tilts, batch 8192, RTX PRO 6000 (idle GPU pinned) s/epoch
torch path (two runs) 2.245 / 2.267
fused dispatch (two runs) 1.481 / 1.601

~1.45x end-to-end per epoch, with final losses in family (1.04–1.08e-3 after 4 epochs on both paths).

Dependencies / conflicts

Evidence

  • New TestFusedDispatchParity in tests/ml/test_kplanes.py: fused-vs-torch forward parity (rtol 1e-4) and gradient parity for points + all grids (kernel toggled via monkeypatch), fp64 fallback, CPU path; skipped automatically when CUDA or quantem-cuda is absent. 6/6 pass with the kernel installed.
  • Full suite uv run pytest tests --runslow: 433 passed, 1 skipped; only the pre-existing unrelated ptychography failure.
  • ruff check / format clean on touched files.

Side finding (not addressed here)

KPlanes(use_hybrid_mlp=False) — the constructor default — crashes in get_params() with AttributeError: 'KPlanes' object has no attribute 'sigma_net' because get_params/param_keys reference sigma_net unconditionally while __init__ only creates it under use_hybrid_mlp=True. Hit while benchmarking; flagged for a separate fix.

…n available

The non-tilted KPlanes feature interpolation runs grid_sample + plane
product per multiscale level; grid_sample backward dominates INR epoch
time. quantem-cuda now ships kplanes_fuse, a fused fwd+bwd kernel with
exactly these semantics. Import it guardedly and dispatch to it from
interpolate_ms_features when the inputs qualify (CUDA, fp32, standard
shapes); every other case falls through to the existing torch path, and
installations without quantem-cuda are unaffected.

Parity tests cover fused-vs-torch forward and gradients (kernel toggled
via monkeypatch), non-fp32 fallback, and the CPU path. End-to-end on the
200-cube phantom (29 tilts, RTX PRO 6000, batch 8192): 2.25-2.27 s/epoch
torch -> 1.48-1.60 s/epoch fused, ~1.45x.
@cedriclim1

Copy link
Copy Markdown
Owner Author

Closing without merge: this dispatches to quantem.cuda.core.ml.kplanes_fuse, which quantem-cuda does not export (inventory: kplanes_tilted_fuse, kplanes_tilted_tv_fuse, tv_loss_{iso,sq,l1}_3d). The guard makes it a permanent no-op — the ImportError sets the handle to None and every call takes the torch fallback — and it also bypasses the use_cuda_kernels config policy that #21 standardizes. Revisit if quantem-cuda grows a non-tilted kernel; until then #13's hoist is the live optimization for this path.

@cedriclim1 cedriclim1 closed this Jul 9, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant