Clamp the runtime start in slice and slice_update so the window stays in bounds by kapellirohith · Pull Request #3917 · ml-explore/mlx

kapellirohith · 2026-07-25T18:31:23Z

Summary

mx.slice and mx.slice_update accept the start position as an array (the
runtime form). That start becomes a flat offset

offset = Σ start[i] * strides[axes[i]]

which is applied straight to the data pointer. Nothing range-checks it.
normalize_dynamic_slice_inputs (mlx/ops.cpp:744) validates the count, the
ndim, the dtype and the axes of start — but never its value — and
neither compute_dynamic_offset nor the dynamic copy kernel bounds-checks the
offset it produces:

// mlx/backend/metal/kernels/copy.h:218-229
dst[dst_idx + dst_offset] = src[src_idx + src_offset];

An out-of-range start therefore reads and writes outside the array. It reaches
five entry points — slice, slice_update, both gradients, and the compiled
path — and all three backends carry the same unchecked arithmetic.

This PR clamps the start to [0, operand_dim - window_dim] in the CPU, Metal
and CUDA offset paths, matching XLA's dynamic-slice / dynamic-update-slice,
which is the operation this API mirrors. It also fixes a signedness bug in the
CPU offset path (§4) found while validating the clamp.

In-range starts are unaffected.

1. Reproducer

Deterministic, three lines:

import mlx.core as mx
x = mx.arange(16).reshape(4, 4)
mx.eval(mx.slice(x, mx.array([-1, -1]), (0, 1), (2, 2), stream=mx.cpu))

On 973e27f8 this terminates the process with SIGBUS, 3 runs out of 3.

The write side, on Metal, silent — the stored value is nowhere in the
destination, so the store went outside it:

dst = mx.zeros((8, 8), dtype=mx.int32)          # valid starts are 0..7
out = mx.slice_update(dst, mx.full((1, 1), 7, dtype=mx.int32),
                      mx.array([100, 0]), (0, 1))
mx.eval(out)
print((out == 7).any())   # False

No exception is raised, and Metal shader validation
(MTL_DEBUG_LAYER=1 MTL_SHADER_VALIDATION=1) reports nothing, because the offset
still lands inside the process's own Metal heap — just in a different tensor.

Whether an out-of-range access faults or silently returns data depends on what
happens to be mapped at the computed address. Both outcomes are observed below
(80 faults, all on CPU; Metal never faulted in 311 cases). I mention this so the
crash is not read as the defining symptom — the silent case is the more common
one and the more dangerous.

Illustrative only. The store can be observed landing in another allocation:

victim = mx.full((1 << 18,), 0x5A5A5A, dtype=mx.int32)
dst    = mx.zeros((8, 8), dtype=mx.int32)
mx.eval(victim, dst)
mx.eval(mx.slice_update(dst, mx.full((1, 1), 0x1234, dtype=mx.int32),
                        mx.array([19537, 0]), (0, 1)))
# on one run: victim[712] went 0x5A5A5A -> 0x1234

The specific numbers — start 19537, index 712 — are allocator-dependent and
will not reproduce byte-for-byte elsewhere. Treat this as a demonstration of
what the unclamped offset can reach, not as a test. The reproducible evidence is
the SIGBUS above, the "store lands outside dst" probe, and the audit in §2.

2. Coverage audit

I did not want to rely on a pass count, so the surface is enumerated as 311
independent cases per backend, each run in its own process so a fault cannot
abort the run. The oracle needs no knowledge of the implementation: a dynamic
slice with start S must equal the same slice taken at
clamp(S, 0, dim - window), computed in NumPy — which is exactly the XLA rule.

Cases cross: entry point × geometry × start kind × start dtype × operand layout.

axis of variation	values
entry point	`slice`, `slice_update`, gradient w.r.t. operand, gradient w.r.t. update, `mx.compile`d
geometry	1-D, 2-D all-axes, 2-D partial axes, 3-D partial axes, 4-D all-axes, window == full axis, unit window
start kind	in-range, zero, max edge, one past, −1, +100, −100
start dtype	`int8/16/32/64`, `uint8/16/32/64`
operand layout	contiguous, transposed, strided

Result — before.txt and after.txt hold the raw per-case output:

	pristine `973e27f8`	with patch
CPU	176 / 311 failing (96 wrong results + 80 faults: 47 SIGBUS, 33 SIGSEGV)	0 / 311
Metal	176 / 311 failing (all silent wrong results, 0 faults)	0 / 311

Failures by entry point (both backends): slice 92, slice_update 92,
gradient-w.r.t.-operand 56, gradient-w.r.t.-update 56, compiled 56.

Failures by start kind: one-past 114, −1 98, +100 70, −100 70, and in-range,
zero and max-edge: 0. The in-range cases are the no-regression controls and
pass on both builds.

Two things this audit turned up that I would otherwise have missed:

The defect reaches gradients. DynamicSlice::vjp and
DynamicSliceUpdate::vjp (mlx/primitives.cpp:5202, :5278) rebuild a
slice / slice_update with the same start, so the gradient took the same
out-of-bounds offset as the forward pass. Both gradient entry points fail on
pristine and pass with the patch, and the patch keeps forward and backward
clamping identically, so the gradient still lands on the window the forward
pass used.
My first gradient-w.r.t.-update case was not discriminating. With a
constant cotangent that gradient is all ones wherever the window lands, so it
passed on pristine. It only became load-bearing after weighting the cotangent
by position. Noting it because a reviewer should not read "grad is covered" as
automatically meaning "grad is tested".

Not affected, checked and excluded: slice_update_add/prod/max/min have no
array-start overload (mlx/ops.cpp — only three functions take const array& start), and vmap over the start indices already throws
(mlx/primitives.cpp:5184, :5241).

3. In-tree regression test

test_dynamic_slicing_start_is_clamped in python/tests/test_ops.py, next to
test_dynamic_slicing. It names both streams explicitly rather than relying
on the ambient device, because MLXTestCase only switches device through the
DEVICE environment variable — a default run would otherwise have exercised
the GPU only, while the faults and the signedness fix are on the CPU side. The
gpu leg is guarded by mx.metal.is_available() or mx.cuda.is_available(), so it
covers whichever GPU backend is built (see §11).

It covers, per backend: out-of-range and negative starts on slice; partial
axes; a window covering the whole axis (only valid start is 0); all eight start
dtypes, asserting signed-negative clamps to 0 and large-unsigned clamps to the
far end; slice_update landing at the clamped position; a neighbouring
allocation left bit-identical; and the gradient landing on the clamped window.

Before / after (intree_before.txt, intree_after.txt):

pristine: records subtest failures, then the process dies with SIGBUS
(signal 10) inside the CPU section. Because that kills the runner, the Metal
section never executes, so I ran it standalone on pristine as well — 5 of 5
Metal assertions fail (returns data outside x; the update lands nowhere in
dst).
with patch: passes.

4. Second defect: signed starts read through unsigned pointers

Independent of the clamp, and visible by inspection. The CPU offset dispatch
funnels signed and unsigned dtypes through the same unsigned pointer type
(mlx/backend/cpu/primitives.cpp:56-72 on 973e27f8):

  switch (indices.dtype()) {
    case int8:
    case uint8:
      encoder.dispatch(compute_offset, indices.data<uint8_t>());
      break;
    case int16:
    case uint16:
      encoder.dispatch(compute_offset, indices.data<uint16_t>());
      break;
    case int32:
    case uint32:
      encoder.dispatch(compute_offset, indices.data<uint32_t>());
      break;
    case int64:
    case uint64:
      encoder.dispatch(compute_offset, indices.data<uint64_t>());
      break;

An int32 start of -1 is read as 4294967295, so the offset becomes
4294967295 * stride rather than -stride. That is why the CPU faults on
[-1, -1] instead of reading just before the buffer. It also means a clamp alone
is insufficient: without this fix a negative start clamps to the far end of the
axis instead of to zero. I hit exactly that in the first iteration of this patch,
and the dtype cases in the audit and the in-tree test now pin it.

Metal is unaffected — its generated kernel is typed with the real dtype.

5. Open question for maintainers: negative starts

The static form treats a negative start as counting from the end; this patch
clamps a negative runtime start to zero. Both are defensible and it is a one-line
change either way, so I would rather you pick than have me choose silently.

Measured on 973e27f8, x = arange(16).reshape(4, 4):

expression	static form today	runtime form, before this PR
`x[-1:, -1:]`	`[[15]]` (counts from the end)	SIGBUS on CPU
`x[-100:, :]`	whole array	out-of-range read
`x[99:, :]`	empty, shape `(0,4)`	out-of-range read
`x[2:100, 2:100]`	`[[10,11],[14,15]]`	out-of-range read
`y[-1:, -1:] = 9`	writes the corner	store lands outside `y`
`y[99:, :] = 9`	no-op	store lands outside `y`

So the same logical operation is already safe when the start is a Python value
and unsafe when it is an array, and nothing in the docstrings marks the runtime
form as weaker.

Option A — clamp negatives to 0 (what this patch does). Matches XLA
dynamic-slice, so code ported from JAX behaves identically. Diverges from the
static form's Python semantics.
Option B — wrap negatives (start + dim), then clamp. Consistent with the
static form and with Python indexing. Diverges from XLA.

Either is a strict improvement over the status quo, and neither changes in-range
behaviour. Say the word and I will switch it and update the tests.

6. Root cause

file	line (on `973e27f8`)	what
`mlx/ops.cpp`	744	`normalize_dynamic_slice_inputs` checks count / ndim / dtype / axes of `start`, never its value
`mlx/backend/metal/slicing.cpp`	76-79	`acc += indices[i] * strides[axes[i]]` — offset built unchecked
`mlx/backend/cpu/primitives.cpp`	48-55	same arithmetic on the CPU side
`mlx/backend/cuda/slicing.cpp`	91-100	same arithmetic on the CUDA side
`mlx/backend/cpu/primitives.cpp`	56-72	signed start dtypes dispatched through unsigned pointers (§4)
`mlx/backend/metal/kernels/copy.h`	218-229	`dst[dst_idx + dst_offset] = src[src_idx + src_offset]` — offset applied with no bounds check

7. The fix

Clamp each start index to [0, operand_dim - window_dim] inside
compute_dynamic_offset, where strides and axes are already passed as
constants — so there are no extra kernel launches on the hot path.

mlx/backend/metal/slicing.cpp — clamp in the generated offset kernel; the
comparison is emitted per dtype in the index's own signedness.
mlx/backend/cpu/primitives.cpp — same clamp, plus signed dtypes are now read
through signed pointers (§4).
mlx/backend/cuda/slicing.cpp — the same clamp in the CUDA offset kernel.
See §10 for what is and is not validated locally there.
mlx/backend/gpu/slicing.h, mlx/backend/gpu/primitives.cpp — pass the
per-axis bounds: in.shape(ax) - out.shape(ax) for DynamicSlice,
out.shape(ax) - upd.shape(ax) for DynamicSliceUpdate.

8. Why clamp rather than raise

The obvious alternative is to reject an out-of-range start with an exception.
That is not free here.

The start is a runtime device value. It is not known on the host at graph
construction time. Verified from the code: nothing in mlx/ops.cpp reads an
input array's values while building an op, and in the dynamic-slice path the
start is consumed only on the device (mlx/backend/metal/slicing.cpp launches a
kernel over the indices buffer; the host never dereferences it). Raising from
slice / slice_update would require evaluating and reading back start on
every call — a device→host synchronisation on the exact op whose purpose is to
avoid one. The runtime form exists so a KV-cache append can be issued without
syncing on the current sequence length.

A device-side check would report in the wrong place. Metal offers no
exception mechanism, and MLX has no kernel-authored error channel: the only error
path in the Metal backend is Device::error_ (mlx/backend/metal/device.h:122),
which carries MTLCommandBufferStatusError from a completed command buffer and
is surfaced at synchronize. A kernel could write a flag to a buffer, but the host
would observe it only at the next synchronisation — detached from the call that
caused it, unless you sync eagerly, which reintroduces the cost above. I am not
claiming a device-side error is impossible, only that it is a larger change and
still pays a sync to be timely.

The status quo is not a loud error. Today an out-of-range start is a silent
out-of-bounds access or a process fault — worse than either clamping or throwing.

If you would rather have a validating error despite the sync, the natural place is
an opt-in debug check rather than the default path, and I am happy to add it.

9. Validation

check	result
Coverage audit, 311 cases × 2 backends, one process each	pristine 352 failing (incl. 80 faults) → 0
`test_dynamic_slicing_start_is_clamped` (new)	pristine: SIGBUS in the CPU section, 5/5 Metal assertions fail standalone → passes with patch
Python suite (`python/tests`)	789 passed, 3 skipped
Python suite with `DEVICE=cpu` (mirrors the CI cpu leg)	789 passed, 6 skipped
C++ suite `DEVICE=gpu`	260 cases / 3489 assertions, 0 failed
C++ suite `DEVICE=cpu`	260 cases / 3493 assertions, 0 failed
Differential sweep, in-range starts, GPU vs CPU, random shapes / axes / dtypes	0 mismatches over 500 cases × 2 ops
Perf, `slice_update` KV-cache append (`[1,8,4096,128]` fp16, 300 iters)	before 313.3 / 318.5 / 308.8 µs/iter; after 318.0 / 322.6 / 320.0 µs/iter

Perf is unchanged within run-to-run spread: the clamp is a few integer operations
inside the existing single-thread offset kernel plus one small constant buffer.

Everything in this table is CPU and Metal. The CUDA hunk is covered by §11.

10. Notes

CUDA is fixed here but validated by CI, not by me. I have no CUDA device,
so the CUDA hunk is not compiled or run locally. See §11.
How an out-of-range start arises unintentionally: [BUG] Metal slice_update can use a stale array-valued start across command buffers #3880 reports that on
Metal a dynamic slice_update can read a stale start across a command-buffer
boundary. I reproduced that on 973e27f8 (17 of 20 fresh processes) and
confirmed the offset is exactly the recycled buffer's prior contents — sweeping
a poison value P moved the write to flat index 9P on an 8×8 array, i.e.
Σ start_i·stride_i verbatim. This PR does not fix [BUG] Metal slice_update can use a stale array-valued start across command buffers #3880; it bounds the
damage when a start is out of range for any reason. The two are independent.
The runtime-start slice_update is the preallocate-and-update KV-cache
pattern, where the start comes from sequence length, so an index exceeding the
preallocated capacity is a realistic way to reach this.
Not run: iPhone (A18 Pro / A17 Pro) via mlx-swift. The defect and the fix
live in host-side C++ and in a scalar offset kernel, so A-series behaviour
should match, but that is reasoning rather than a measurement.

11. CUDA: fixed here, validated by CI rather than locally

I have no CUDA device, so the CUDA hunk is not compiled or run on my machine.
Leaving it out was not an option: compute_dynamic_offset is declared once in
mlx/backend/gpu/slicing.h and mlx/backend/gpu/primitives.cpp calls it for
every GPU backend, so adding the bounds argument without updating
mlx/backend/cuda/slicing.cpp would have left the CUDA builds unable to link.

What I did to make the blind change defensible. The CUDA kernel is a JIT
template whose only include is mlx/backend/cuda/device/utils.cuh, so I checked
each thing it relies on against what that header actually provides:

max_starts mirrors the existing axes parameter exactly — same
const __grid_constant__ cuda::std::array<int, NIDX> type, same
args.append(...) overload (std::vector<int> → SmallVector<int> →
pointer, jit_module.h:49-52), same length NIDX. Argument order in
KernelArgs matches the kernel signature.
if constexpr is already used in NVRTC-visible device headers
(device/binary_ops.cuh:19), so C++17 in this context is not an assumption.
The signedness test uses cuda::std::numeric_limits<T>::is_signed rather than
is_signed_v, because utils.cuh includes <cuda/std/limits> directly
while <cuda/std/type_traits> only arrives transitively. T is always a real
fixed-width integer type here (dtype_to_cuda_type, backend/cuda/utils.cpp:32-47),
and the start dtype is validated as integer at the op level.
I extracted the exact kernel source string that is handed to NVRTC and
compiled it as host C++17 with -Wall -Wextra against stubs for
__global__, __grid_constant__, cuda::std::array,
cuda::std::numeric_limits and Strides, instantiated for all eight integer
start dtypes and NIDX of 1, 2 and 4: clean, no warnings. Running it then gave
the correct clamped offset on 12 probes, including a uint64 start of 2^64-1
(clamps to the far end, not to zero). That covers syntax, types, the
if constexpr, and the arithmetic; it does not cover NVRTC itself or the
__grid_constant__ argument passing.
The kernel source is a R"( ... )" raw literal, so I checked the added lines
introduce no )" sequence that would terminate it early.

CUDA has no signedness bug of its own — unlike the CPU path, its kernel is
templated on the real dtype — so only the clamp was needed there.

What CI will prove. From .github/workflows/build_and_test.yml, the
build_and_test matrix builds cuda-12.6, cuda-12.9 and cuda-13.0 on Linux
for both architectures, and the test step runs when
matrix.os == 'Linux' && (matrix.toolkit == 'cpu' || matrix.arch == 'x86_64')
(lines 71-72). So:

Linux aarch64 + CUDA: compiles and links the CUDA hunk, no tests. This
alone catches the link breakage described above and any host-side compile
error.
Linux x86_64 + CUDA (×3 toolkits): runs on gpu-t4-4-core, a real T4. The
test-linux action probes for a GPU and, when present, runs the Python suite
with DEVICE=gpu and the C++ suite with DEVICE=gpu. That is where the new
test actually exercises the CUDA offset kernel, and where an NVRTC compile
error in the JIT source would surface — NVRTC compiles at runtime, so a
malformed kernel fails the test, not the build.
macOS: mac_build plus a self-hosted mac_test, which is the Metal coverage I
already ran locally.

What CI will not prove. Nothing checks the CUDA path before it runs, so if
the JIT source has a typo it shows up as a failing test rather than as a build
error, and I would need a maintainer's run to see it. If that happens the fix is
confined to one kernel body in mlx/backend/cuda/slicing.cpp and I will turn it
around quickly.

One correction worth flagging, since it changes what CI covers: I originally
guarded the gpu leg of the new test with mx.metal.is_available(), which is
false on a CUDA build — the gpu leg would have been silently skipped on the
T4 runners, so CI would have gone green whether or not CUDA was fixed. The guard
is now mx.metal.is_available() or mx.cuda.is_available(), matching how the rest
of python/tests guards CUDA, so the CUDA kernel is genuinely exercised.

Files

0001-clamp-runtime-start.patch — the commit (6 files, +198/-19, includes the
regression test), parented on 973e27f8
case.py, drive.py — the coverage audit (311 cases per backend, one process each)
before.txt, after.txt — raw per-case audit output on pristine and with the patch
intree_before.txt, intree_after.txt — the in-tree test on pristine and with the patch
repro_probe.py, evidence_before.txt, evidence_after.txt — the earlier
18-probe read/write matrix, kept as independent corroboration of §1

slice and slice_update take the start position as an array. That start is turned into a flat offset of sum(start[i] * strides[axes[i]]) and applied straight to the data pointer, and nothing range checks it along the way: normalize_dynamic_slice_inputs validates the count, ndim, dtype and axes of the start but never its value, and neither compute_dynamic_offset nor the dynamic copy kernel bounds check the offset they produce. An out of range start therefore reads and writes outside the array, silently on the GPU backends and with a fault on the CPU. It reaches slice, slice_update, both vjps and the compiled path. Clamp each start index to [0, operand_dim - window_dim] where the offset is computed, which is what XLA's dynamic-slice does. The bounds go in next to the strides and axes that are already passed as constants, so this costs no extra kernel launches. In range starts are unaffected. The CPU dispatch also read signed start dtypes through unsigned pointers, so an int32 start of -1 arrived as 4294967295. Read each dtype through its own pointer type, and do the comparison in the index's own signedness so that neither a negative signed start nor a large unsigned one clamps to the wrong end of the axis.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Clamp the runtime start in slice and slice_update so the window stays in bounds#3917

Clamp the runtime start in slice and slice_update so the window stays in bounds#3917
kapellirohith wants to merge 1 commit into
ml-explore:mainfrom
kapellirohith:slice-clamp

kapellirohith commented Jul 25, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

kapellirohith commented Jul 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

1. Reproducer

2. Coverage audit

3. In-tree regression test

4. Second defect: signed starts read through unsigned pointers

5. Open question for maintainers: negative starts

6. Root cause

7. The fix

8. Why clamp rather than raise

9. Validation

10. Notes

11. CUDA: fixed here, validated by CI rather than locally

Files

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

kapellirohith commented Jul 25, 2026 •

edited

Loading