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

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
aarmey pushed a commit that referenced this pull request Sep 18, 2026
…ory metrics (#56)

* Assert memory ceilings with pytest-memray instead of tracemalloc

The two memory assertions in the suite reset `tracemalloc`'s peak around the
operation and compared the delta to a hand-computed bound. That works, but
the bound was rebuilt inline at each site, the setup had to be fenced off by
hand, and the number it compared against moved with the runner's core count:
the accumulator is `nthreads * n_minor * width * 8` bytes, so the same
assertion meant something different on a 4-core runner than on a 96-core one.

`pytest-memray` states the same property as a mark. It is what anndata uses
for this, it intercepts the allocator rather than Python's allocation hooks,
and a ceiling is enforced by having the plugin installed -- no `--memray`
flag, so CI needs no change.

Three conventions keep the ceilings meaningful:

  - `pinned_threads` (new, in conftest) fixes numba's thread count, so the
    expected allocation is a property of the code and not of the machine.
    Measured at 4 threads: 20 KB for a minor-axis sum, 40 KB for the extrema
    kernels, 132 KB for the misaligned matmul.
  - Inputs and JIT warm-up move into module-scoped fixtures. A mark measures
    the test body only -- verified -- so the array under test no longer has
    to be fenced off by resetting a counter.
  - The ceiling rides on each `pytest.param`, never on `request.applymarker`.
    The plugin reads the marker at collection time, so a marker applied from
    the body is silently ignored and the test asserts nothing. Every ceiling
    here was checked by tightening it to 1 KB and confirming it fails.

Adds one `limit_leaks` test: a cache that grows with use is invisible to a
ceiling, since each pass alone stays well under it, and that is the shape of
the bug d286cf3 fixed. Confirmed it catches a deliberately retained result.

`benchmarks/` keeps `tracemalloc`: it runs outside pytest, where marks do not
apply, and it records numbers rather than bounding them. Both READMEs now say
which tool owns which question, including that neither sees RSS -- so neither
catches the allocator fragmentation behind #49's 140 GB, which is diagnosable
but not reliably gateable.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* Retire the benchmark suite's memory metrics in favour of the ceilings

With `pytest-memray` asserting the memory properties directly, recording them
here as well left two mechanisms for one claim, and the weaker one gated.

The claims were never about a number. "A minor-axis reduction must not
allocate anything nnz-sized" either holds or it does not; a recorded figure
shows it drifting only after the fact, against a baseline that has to be
re-recorded whenever the bound legitimately moves -- and the checked-in
ceilings were already stale, generous figures taken before the memory fixes
landed, with a table in the README tracking how far off they were. Worse, the
number was never portable: `peak_alloc_mb` measured a thread-local
accumulator sized by the runner's core count, so the same code produced a
different figure on every machine. A ceiling with a pinned thread count says
the thing that is actually meant.

Removed: the five memory-only cases, `peak_alloc_mb_view` and its companions
from the 30 normalized-view cases, `matmul_peak_alloc_mb` from the slow case,
and `peak_alloc_mb` from the harness. 24 gated cases remain.

Nothing loses coverage. Two gaps were filled first rather than dropped:
`minor_selection_peak_mb` had no test counterpart, so
`test_minor_axis_selection_allocates_no_nnz_sized_scratch` covers it -- with a
ceiling above the output, since a selection's result legitimately grows with
what was selected, and below output-plus-an-nnz-sized-temporary. And the 30
per-recipe view cases became a parametrization of the existing matmul ceiling
over every recipe; measured, all five allocate the same 132 KB.

Layout size and throughput stay. `bytes_per_nonzero` is the compression
format's core claim and has no other guard, and the scipy-relative timings
have no equivalent in the test suite. Only the half memray replaces is gone.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
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