Dispatch interpolate_ms_features to the fused quantem-cuda kernel - #16
Closed
cedriclim1 wants to merge 1 commit into
Closed
Dispatch interpolate_ms_features to the fused quantem-cuda kernel#16cedriclim1 wants to merge 1 commit into
cedriclim1 wants to merge 1 commit into
Conversation
…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.
This was referenced Jun 11, 2026
Owner
Author
|
Closing without merge: this dispatches to |
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
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.
What
interpolate_ms_features(non-tilted KPlanes) now dispatches toquantem.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_samplebackward 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:~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
TestFusedDispatchParityintests/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.uv run pytest tests --runslow: 433 passed, 1 skipped; only the pre-existing unrelated ptychography failure.Side finding (not addressed here)
KPlanes(use_hybrid_mlp=False)— the constructor default — crashes inget_params()withAttributeError: 'KPlanes' object has no attribute 'sigma_net'becauseget_params/param_keysreferencesigma_netunconditionally while__init__only creates it underuse_hybrid_mlp=True. Hit while benchmarking; flagged for a separate fix.