Skip to content

Fix the CPU scan over a size one axis with a padded stride - #4139

Open
kapellirohith wants to merge 1 commit into
ml-explore:mainfrom
kapellirohith:cpu-scan-size-one-axis
Open

Fix the CPU scan over a size one axis with a padded stride#4139
kapellirohith wants to merge 1 commit into
ml-explore:mainfrom
kapellirohith:cpu-scan-size-one-axis

Conversation

@kapellirohith

@kapellirohith kapellirohith commented Aug 10, 2026

Copy link
Copy Markdown
Contributor

Problem

A scan over a size one axis whose stride is padded leaves the output buffer
completely unwritten on the CPU, so cumsum and friends return whatever the
allocator last put there.

Minimal reproducer (M3 Pro, macOS 26.x, main @ 3f2e4a3):

import mlx.core as mx
import numpy as np

# fill a buffer of exactly the output size with a marker, then release it
tmp = mx.full((7,), 12345.0, dtype=mx.float32)
mx.eval(tmp)
del tmp

base = mx.arange(1, 11, dtype=mx.float32).reshape(1, 10)
mx.eval(base)
a = base[:, 3:]                      # shape (1, 7), strides (10, 1)
mx.eval(a)

out = mx.cumsum(a, axis=0, stream=mx.cpu)
print(np.array(out))
# before: [12345. 12345. 12345. 12345. 12345. 12345. 12345.]
# after:  [    4.     5.     6.     7.     8.     9.    10.]

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:

variant result
cumsum(a, axis=0, stream=mx.cpu) wrong, returns the recycled buffer
cumsum(a, axis=0, stream=mx.gpu) correct
base[:, 0:] instead of base[:, 3:] correct
axis=1, the size 7 axis correct
base itself, no slice correct

It 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_op in mlx/backend/cpu/scan.cpp picks the strided path whenever the
scanned axis does not have stride one, and computes the row count as

in.size() / in.shape(axis) / in.strides()[axis]

row_contiguous deliberately exempts size one axes: check_contiguity requires
strides[i] == prod(shape[i+1:]) only when shape[i] != 1, so a size one axis
may 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 the
count is 7 / 1 / 10, which floors to zero, and the
for (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 init for exclusive.

That relies on a row contiguous array being densely packed in logical order over
in.size() elements. It is: for every axis with shape[i] > 1,
check_contiguity forces strides[i] == prod(shape[i+1:]), which is the mixed
radix 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_scans slices only the outputs it compares, never the input, so every
existing 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_gpu sizes the output from in.data_size() and
then bounds the kernel writes by in.strides(). That is a separate defect in a
separate 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 code
path it is testing and away from the Metal one. I instrumented Scan::eval_gpu
to 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.

check before after
scan sweep: 4 ops x reverse x inclusive x 10 dtypes x every generator of a size one axis 1312 of 3344 wrong 0 of 3344, both devices
as_strided with gaps and absurd strides, cross backend cpu vs gpu differential, donation, multiple size one axes, compile cold and warm n/a 0 of 95 per device
exhaustive row contiguous search, 1860 candidates n/a 0 of 4560
lazy eval, subgraph then parent, eval twice, async_eval, two cpu streams on one buffer, 50 run bit identical determinism, export round trip n/a 0 of 6 per device
the new test RED 5/5 on both DEVICE=cpu and DEVICE=gpu OK, standalone and inside the full test_ops module
200 consecutive runs of the new test n/a 0 failures per device
python suite, GPU and mx.set_default_device(mx.cpu) n/a 832 tests, pre-existing test_fft_too_large on cpu only
C++ suite, both devices n/a 266 cases
ASAN and UBSAN, CPU backend, C++ suite n/a 0 AddressSanitizer errors, no finding in cpu/scan.cpp
pre-commit n/a clean

Because 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 than
from the stride, so this defect does not apply to it. That is from reading the
source, I have no NVIDIA GPU.

Checklist

Put an x in the boxes that apply.

  • I have read the CONTRIBUTING document
  • I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (if needed)

@kapellirohith
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
kapellirohith force-pushed the cpu-scan-size-one-axis branch from 93fa931 to 9eded5b Compare August 13, 2026 09:08
@kapellirohith
kapellirohith marked this pull request as ready for review August 13, 2026 09:08
@kapellirohith
kapellirohith marked this pull request as draft August 13, 2026 09:12
@kapellirohith
kapellirohith marked this pull request as ready for review August 13, 2026 09:50
@zcbenz zcbenz added the await verification This pull request is non-trivial and requires a human expert to verify its correctness. label Aug 13, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

await verification This pull request is non-trivial and requires a human expert to verify its correctness.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants