Skip to content

Apply centered-patches optimization to the GPU path - #1

Merged
YQ-Wang merged 2 commits into
mainfrom
centered-matrix
Jun 27, 2026
Merged

Apply centered-patches optimization to the GPU path#1
YQ-Wang merged 2 commits into
mainfrom
centered-matrix

Conversation

@YQ-Wang

@YQ-Wang YQ-Wang commented Jun 27, 2026

Copy link
Copy Markdown
Owner

What

Brings the GPU path in line with the CPU path's "centered patches" formulation and refreshes the benchmarks.

  • CPU path (existing, commit 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).
  • GPU path (this PR) now does the same: the CUDA sparse tensor is built from _centered_patches(...) and the per-batch diff_ik.sub_(k_nn * exp_batch_gpu) is removed.
  • Removed the now-redundant call-site comments (the _centered_patches docstring already explains the rationale).
  • Refreshed the README Reference performance and graph-reuse benchmarks.

Why

The change is algebraically exact:

(patches_raw - k_nn*I) @ exp  ==  patches_raw @ exp - k_nn*exp

patches_raw has a zero diagonal (self is excluded from the k-NN), so centering just sets the diagonal to -k_nn, matching upstream CastleLi/STORM exactly. Both approx and exact GPU 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)

  • Tests: all 33 pass, including the three CPU↔GPU equivalence tests (test_storm_sparse_gpu_equivalence, test_storm_exact_gpu_equivalence, test_cached_gpu_results_match).
  • Numerically identical: CPU↔GPU corr = 1.000000, max p-value/effect-size diff < 5e-5, identical p < 0.05 calls — unchanged before vs. after.
  • Timing-neutral: GPU time within run-to-run noise (the removed sub_ was cheap next to the SpMM + float64 column reductions).
  • Memory-neutral: peak GPU VRAM changes by sub-MB, scaling linearly with cell count (the one extra -k_nn diagonal nonzero per row); the removed per-batch temporary wasn't the high-water mark (the cuSPARSE SpMM is).

Reviewer notes

  • The PR also includes the prior CPU-centering commit db1ff4d already on this branch.
  • README benchmark numbers were re-measured (median of five) and the graph-reuse paragraph updated accordingly.

YQ-Wang and others added 2 commits June 26, 2026 22:43
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>
@YQ-Wang
YQ-Wang merged commit f8541d2 into main Jun 27, 2026
5 checks passed
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