Description

mx.pad produces incorrect values on Metal without raising an exception. With an output row stride of 16, the first tested bad row is 134217728, whose destination element offset is exactly 2^31. The same operation returns correct results on CPU.

This reproducer requires no model, images, inference server, cache, or network. It allocates an output array just over 2 GiB, so please run it with sufficient available memory.

Environment tested

• MLX 0.32.4 and mlx-metal 0.32.4, stock binaries bundled with oMLX 0.7.1
• Python 3.11
• macOS 27.0.1 (26A434), arm64
• Apple M5 Max, 128 GiB unified memory
• No locally built MLX library or proposed C++ fix was installed

Minimal reproduction

import mlx.core as mx
import numpy as np

rows = 2**27 + 1
for device in (mx.gpu, mx.cpu):
    with mx.stream(device):
        x = mx.broadcast_to(mx.array([[1, 2, 3]], dtype=mx.uint8), (rows, 3))
        y = mx.pad(x, ((0, 0), (0, 13)))
        mx.eval(y)
        for n in (0, rows - 2, rows - 1):
            print(device, n, np.array(y[n:n + 1]).reshape(-1).tolist())
        del x, y
    mx.clear_cache()

Expected and observed results

Every row should be [1, 2, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0].

There is no exception or allocation failure. The input values in the last tested Metal row are silently lost.

Suspected source-level cause

In pad_gpu at v0.32.4, the destination view is created using:

array out_slice(in.shape(), out.dtype(), nullptr, {});
out_slice.copy_shared_buffer(
    out, out.strides(), out.flags(), out_slice.size(), data_offset);

The view retains the padded output's strides but reports the input's logical element count as its data size. In the reproducer, the logical size is 134217729 * 3 = 402653187 elements, while the last input value must be written at destination offset 134217728 * 16 + 2 = 2147483650.

copy.cpp selects the large-index copy kernel based on the reported input/output data_size(). The reported size can therefore select int32 indexing even when the true destination offsets exceed INT32_MAX. This is consistent with the observed boundary, but I have not built or verified a native C++ fix.

Related: #4548 / #4540 describe/fix a similar span-reporting issue for split. This report concerns the pad destination view and still reproduces on 0.32.4, which already includes the split fix.

Could the padding destination view preserve its true strided span (and valid view flags), or could copy-kernel dispatch check actual address ranges? A regression test straddling the above offset would help prevent silent data corruption.

Downstream impact (secondary evidence)

This was first observed during multi-image Qwen vision encoding: patch-embedding Conv3d internally pads RGB channels from 3 to 16. Large batches produced corrupted tail-image features before attention or any cache write. Independent constant-input Conv3d and then the pad-only reproducer narrowed the failure down to padding.

As a diagnostic workaround only, splitting the independent patch-embedding convolution into smaller chunks and concatenating its outputs preserved all final per-image feature hashes against single-image references: 120 images at each of two resolutions, 240/240 matches. The rest of vision encoding still consumed the full image sequence. This is not a claim that a native MLX fix has been validated.

The pad-only reproduction above should be sufficient to investigate without any private media or model download.

repro.py