Replace chunked-transpose matmul fallback with a native minor-axis kernel - #49
Merged
Conversation
…rnel The misaligned direction of a normalized view's matmul (self@B for VCSC, B@self for VCSR) previously regrouped the array into the other VCS format a chunk of major slices at a time (_chunk_bounds/_transpose_major), caching a full opposite-format copy (_dual_arr) when the whole array fit one chunk's budget. At scale (~2.3B nonzeros, 1.3M x 12.2K), this needed ~1,150 chunks per pass, repeated across every power iteration and every rank/trial in a BiCV sweep -- tens of thousands of variably-sized alloc/free cycles that fragmented the allocator and drove RSS to ~140 GB on a shared machine, even though no single chunk's live memory was large. Replace it with a direct minor-axis kernel that walks the array's own storage as-is: each thread gets a private, full-output-sized accumulator and scatters into it while owning a disjoint range of major slices, summed across threads at the end -- the same pattern _ops.py already uses for minor_sums/minor_counts/minor_extrema. accumulator_threads caps the thread count to a fixed byte budget, so the accumulator block is always nthreads * output_dim * width * 8 bytes: bounded, independent of nnz, and allocated once per call instead of thousands of times. For a huge output axis (e.g. millions of cells) that caps down to a single thread, trading parallelism for a hard memory bound rather than the previous unbounded chunk churn. Removes _chunk_bounds/_aligned_source/_build_dual/_dual_arr and the chunked-transpose test suite; adds tests/test_minor_axis_matmul.py covering correctness against a dense reference, thread-count invariance, that no second copy of the array is ever built, and that peak memory stays bounded by the accumulator budget rather than nnz. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This was referenced Sep 13, 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
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.
Summary
The misaligned direction of a normalized view's matmul (
self @ BforVCSCArrayNormalized,B @ selfforVCSRArrayNormalized) previouslyregrouped the array into the other VCS format a chunk of major slices at a
time (
_chunk_bounds/_transpose_major), caching a full opposite-formatcopy (
_dual_arr) whenever the whole array's regroup fit one chunk'sbudget.
At the scale this was built for (BAL-Pf2's ~2.3 billion-nonzero, 1.3M x
12.2K matrix), that fallback needed ~1,150 chunks per pass, repeated across
every power iteration and every rank/trial of a BiCV sweep -- tens of
thousands of variably-sized alloc/free cycles. Each chunk was individually
bounded and freed, but that many alloc/free cycles of varying size
fragmented the allocator (glibc malloc not returning freed memory to the
OS), driving RSS to ~140 GB on a shared machine even though the live
working set at any instant was small.
This replaces the chunked/dual-cache fallback with a direct minor-axis
kernel that walks the array's own storage as-is -- no regrouping, no second
copy of the array ever built. Each thread gets a private, full-output-sized
accumulator and scatters into it while owning a disjoint contiguous range of
major slices; the private copies are summed once at the end. This mirrors
the pattern
vsparse._opsalready uses forminor_sums/minor_counts/minor_extrema(accumulator_threads), reused here rather thanreimplemented.
accumulator_threadscaps the thread count so the accumulator block neverexceeds a small fixed budget (
nthreads * output_dim * width * 8bytes) --bounded and independent of nnz, allocated once per call instead of
thousands of times. For a huge output axis (e.g. millions of cells) it caps
down to a single thread, trading parallelism for a hard memory bound rather
than unbounded chunk churn -- the right trade at that scale.
Changes
_vcs_matmul.py: new_vcsc_matmul_delta_minor/_vcsr_rmatmul_delta_minorkernels; removed
_chunk_bounds,_aligned_source,_build_dual,_CHUNK_BUDGET_BYTES,_TRANSPOSE_BYTES_PER_NNZ._vcs_norm.py: removed_dual_arr(nothing builds or reads it anymore)._norm_common.py: updated docstrings that referenced_dual_arr.tests/test_chunked_transpose.py; addedtests/test_minor_axis_matmul.py(correctness vs. dense reference,thread-count invariance, no second array copy is built, peak memory
bounded by the accumulator budget rather than nnz, and that the budget
degrades to 1 thread for a huge output axis).
tests/test_vcs_norm.py/test_vcs_norm_recipes.pywhere theyreferenced the removed
_dual_arrcaching behavior.Test plan
uv run pytest-- 1199 passed, 51 skipped, 3 pre-existing failuresunrelated to this change (confirmed identical on
mainwithout thisdiff: two
test_property_normalization.py::test_recipe_matches_referenceedge cases and one
test_property_norm_stats.py::test_norm_sq_matches_dense_referencecase, all on a degenerate all-constant-column matrix)
uv run ruff check ./uv run ruff format --check .uv run ty check src/to 50,000 x 2,000 (15M nonzeros), and thread-budget sanity-checked at
BAL-Pf2's actual shape (12,210 genes, rank/width up to 100): caps to
3-13 threads depending on width, accumulator block ~15-59 MB, well
under the 64 MiB budget.
Motivation
This was blocking a downstream BAL-Pf2 BiCV rank-selection run that OOM'd a
shared machine (~140 GB RSS) using the chunked-transpose fallback. See
#32 for the PR that introduced that fallback.