Fix the CPU scan over a size one axis with a padded stride - #4139
Open
kapellirohith wants to merge 1 commit into
Open
Fix the CPU scan over a size one axis with a padded stride#4139kapellirohith wants to merge 1 commit into
kapellirohith wants to merge 1 commit into
Conversation
kapellirohith
marked this pull request as draft
August 10, 2026 16:39
row_contiguous exempts size one axes, so such an axis can carry any stride and the array is not copied. The strided scan then takes its row count from that stride, size / shape[axis] / stride, which floors to zero and leaves the output unwritten. A padded stride is only reachable on a size one axis, where the scan is elementwise, so send that case to the contiguous path.
kapellirohith
force-pushed
the
cpu-scan-size-one-axis
branch
from
August 13, 2026 09:08
93fa931 to
9eded5b
Compare
kapellirohith
marked this pull request as ready for review
August 13, 2026 09:08
kapellirohith
marked this pull request as draft
August 13, 2026 09:12
kapellirohith
marked this pull request as ready for review
August 13, 2026 09:50
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.
Problem
A scan over a size one axis whose stride is padded leaves the output buffer
completely unwritten on the CPU, so
cumsumand friends return whatever theallocator last put there.
Minimal reproducer (M3 Pro, macOS 26.x, main @ 3f2e4a3):
The axis being scanned has size one, so an inclusive scan is the identity and
the answer is just the input. Before this change the result is the marker
pattern, byte for byte. That is the positive identification: the output is not
miscomputed, it is never written at all, and what comes back is the previous
occupant of the allocation. Without the marker it is usually zeros, which is
easy to mistake for a real answer.
What makes it fire and what makes it stop, all on main @ 3f2e4a3:
cumsum(a, axis=0, stream=mx.cpu)cumsum(a, axis=0, stream=mx.gpu)base[:, 0:]instead ofbase[:, 3:]axis=1, the size 7 axisbaseitself, no sliceIt needs the CPU backend, a scanned axis of size one, and a stride on that axis
that is larger than the packed extent, which a sliced view gives you.
Mechanism
scan_opinmlx/backend/cpu/scan.cpppicks the strided path whenever thescanned axis does not have stride one, and computes the row count as
row_contiguousdeliberately exempts size one axes:check_contiguityrequiresstrides[i] == prod(shape[i+1:])only whenshape[i] != 1, so a size one axismay carry any stride and the array is still flagged row contiguous and is not
copied. In the reproducer the shape is
(1, 7)with strides(10, 1), so thecount is
7 / 1 / 10, which floors to zero, and thefor (int i = 0; i < count; i++)body never executes.The fix sends that case to the contiguous path instead, which is exactly right
rather than merely safe: a padded stride is only reachable on a size one axis,
and a scan over a size one axis is elementwise, which is what the contiguous
path computes when its stride argument is one. No copy is introduced, and both
forms are covered, since that path with stride one writes the input for
inclusive and
initfor exclusive.That relies on a row contiguous array being densely packed in logical order over
in.size()elements. It is: for every axis withshape[i] > 1,check_contiguityforcesstrides[i] == prod(shape[i+1:]), which is the mixedradix row major encoding over exactly those axes, and the size one axes
contribute index zero regardless of their stride. Enumerating every shape and
stride combination over 2-D and 3-D with dims in {1,2,3} that satisfies that
predicate with at least one size one axis gives 1860 candidates and 4560 scans,
with no counterexample.
Introduced by c423074 "redesign for faster cpu/gpu synch (#1869)"
(2025-03-06).
Why existing tests missed it
No scan test scans an axis of size one, and none scans a sliced view.
test_scansslices only the outputs it compares, never the input, so everyexisting case has stride one on the scanned axis and takes the contiguous path.
Note on the GPU backend
The same input class also makes the Metal strided scan write past its output
allocation, because
Scan::eval_gpusizes the output fromin.data_size()andthen bounds the kernel writes by
in.strides(). That is a separate defect in aseparate file and I will send it separately. It is why the test here pins the
scan to the CPU stream with
stream=mx.cpu: that keeps the test on the codepath it is testing and away from the Metal one. I instrumented
Scan::eval_gputo print every dispatch and confirmed the test issues zero Metal scan dispatches
on the GPU leg, and that it disturbs no canary tensor across ten in-module runs,
on a build carrying this fix but not the Metal one.
Validation
M3 Pro, macOS 26.x, against main @ 3f2e4a3.
as_stridedwith gaps and absurd strides, cross backend cpu vs gpu differential, donation, multiple size one axes, compile cold and warmasync_eval, two cpu streams on one buffer, 50 run bit identical determinism, export round tripDEVICE=cpuandDEVICE=gputest_opsmodulemx.set_default_device(mx.cpu)test_fft_too_largeon cpu onlycpu/scan.cppBecause the test pins the scan to the CPU stream it goes red on both CI legs
rather than only the cpu one, so the GPU leg validates this fix too instead of
skipping it.
Not verified locally: CUDA, Linux, Windows. The CUDA scan shares the Metal
structure, not this one: it takes its row count from
data_size()rather thanfrom the stride, so this defect does not apply to it. That is from reading the
source, I have no NVIDIA GPU.
Checklist
Put an
xin the boxes that apply.pre-commit run --all-filesto format my code / installed pre-commit prior to committing changes