Replace chunked-transpose matmul fallback with a native minor-axis kernel - #49
Merged
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 was referenced Sep 18, 2026
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>
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.