Apply centered-patches optimization to the GPU path - #1
Merged
Conversation
The CPU path already folds upstream's centroid term into the neighbor matrix (PatchesCells = PatchesCells_Raw - K_NN * Diagonal) so each feature batch is a single sparse matmul. Mirror that on the GPU path: build the CUDA sparse tensor from _centered_patches() and drop the per-batch diff_ik.sub_(k_nn * exp_batch_gpu). This is algebraically identical -- (patches_raw - k_nn*I) @ exp == patches_raw @ exp - k_nn*exp -- so results are unchanged (CPU/GPU equivalence tests pass; max p-value/effect-size diff < 5e-5) while unifying the CPU and GPU paths onto the same upstream formulation. GPU timing and peak VRAM are unchanged within measurement noise. Also remove the now-redundant centering call-site comments (the _centered_patches docstring already explains the rationale) and refresh the README reference-performance and graph-reuse benchmarks (RTX 5070 Ti, median of five). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
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
Brings the GPU path in line with the CPU path's "centered patches" formulation and refreshes the benchmarks.
db1ff4d) folds upstream's centroid term into the neighbor matrix via_centered_patches()—PatchesCells = PatchesCells_Raw - K_NN * Diagonal— so each feature batch is a single sparse matmul (PatchesCells %*% ExpMat)._centered_patches(...)and the per-batchdiff_ik.sub_(k_nn * exp_batch_gpu)is removed._centered_patchesdocstring already explains the rationale).Why
The change is algebraically exact:
patches_rawhas a zero diagonal (self is excluded from the k-NN), so centering just sets the diagonal to-k_nn, matching upstream CastleLi/STORM exactly. BothapproxandexactGPU modes share this matmul, so both are covered. The result is one unified formulation across CPU, GPU, and upstream — and one fewer per-batch temporary/kernel on GPU.Verification (NVIDIA RTX 5070 Ti, torch 2.10 + CUDA 12.8)
test_storm_sparse_gpu_equivalence,test_storm_exact_gpu_equivalence,test_cached_gpu_results_match).corr = 1.000000, max p-value/effect-size diff< 5e-5, identicalp < 0.05calls — unchanged before vs. after.sub_was cheap next to the SpMM + float64 column reductions).-k_nndiagonal nonzero per row); the removed per-batch temporary wasn't the high-water mark (the cuSPARSE SpMM is).Reviewer notes
db1ff4dalready on this branch.