Skip to content

Replace chunked-transpose matmul fallback with a native minor-axis kernel - #49

Merged
aarmey merged 1 commit into
mainfrom
minor-axis-matmul-kernel
Sep 13, 2026
Merged

Replace chunked-transpose matmul fallback with a native minor-axis kernel#49
aarmey merged 1 commit into
mainfrom
minor-axis-matmul-kernel

Conversation

@aarmey

@aarmey aarmey commented Sep 13, 2026

Copy link
Copy Markdown
Member

Summary

The misaligned direction of a normalized view's matmul (self @ B for
VCSCArrayNormalized, B @ self for VCSRArrayNormalized) 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) whenever the whole array's regroup fit one chunk's
budget.

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._ops already uses for minor_sums/minor_counts/
minor_extrema (accumulator_threads), reused here rather than
reimplemented.

accumulator_threads caps the thread count so the accumulator block never
exceeds a small fixed budget (nthreads * output_dim * width * 8 bytes) --
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_minor
    kernels; 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.
  • Removed tests/test_chunked_transpose.py; added
    tests/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).
  • Updated tests/test_vcs_norm.py/test_vcs_norm_recipes.py where they
    referenced the removed _dual_arr caching behavior.

Test plan

  • uv run pytest -- 1199 passed, 51 skipped, 3 pre-existing failures
    unrelated to this change (confirmed identical on main without this
    diff: two test_property_normalization.py::test_recipe_matches_reference
    edge cases and one test_property_norm_stats.py::test_norm_sq_matches_dense_reference
    case, all on a degenerate all-constant-column matrix)
  • uv run ruff check . / uv run ruff format --check .
  • uv run ty check src/
  • Manual benchmark: correctness + timing on synthetic VCSR matrices up
    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.

…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>
@aarmey
aarmey merged commit 0335f99 into main Sep 13, 2026
6 checks passed
@aarmey
aarmey deleted the minor-axis-matmul-kernel branch September 13, 2026 20:31
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