Summary
On the iOS GPU (Apple A-series Metal), MLX.conv1d (and convTransposed1d) return incorrect
results when the output length along the convolution axis exceeds 32768 (2¹⁵). Output positions
before roughly (length − 32768) collapse to a near-constant value; only the last ~32768
positions are correct. The same code is correct on the macOS GPU and on CPU.
It looks like a 2¹⁵ addressing/indexing limit in the Metal convolution kernel (block or threadgroup
offset overflowing a 16-bit index).
Impact
Found while running the Kokoro neural TTS vocoder (an iSTFTNet HiFi-GAN-style decoder) on device:
long sentences produced a sustained tone/"beep" because the vocoder's final conv layers emit a
constant spectrogram region. Anything that runs a 1-D conv over a long sequence on iOS (audio
vocoders, long time-series) is affected. The 32768 boundary is well within reach — ~1.4 s of 24 kHz
audio at the post-upsampling resolution.
Environment
- mlx-swift 0.30.2 (also reproduces via the bundled core
mlx)
- Device: iPhone (Apple A-series GPU), iOS 18.x. Does not reproduce on macOS (Apple Silicon) GPU.
- Channels-last layout
[batch, length, channels], float32.
Minimal reproducer (self-contained, no reference output needed)
A 1-D convolution is translation-equivariant, so the first N outputs must not depend on whether
the input is N long or much longer (away from the padding edges). On the iOS GPU they differ:
import MLX
import MLXRandom
// Channels-last input, well past the 32768 output boundary.
let L = 70_000
let x = MLXRandom.normal([1, L, 8])
let w = MLXRandom.normal([4, 3, 8]) // [outChannels, kernelSize, inChannels]
let yFull = conv1d(x, w, stride: 1, padding: 1) // [1, 70000, 4]
let ySub = conv1d(x[0..., 0 ..< 40_000, 0...], w, stride: 1, padding: 1) // [1, 40000, 4]
// Compare an interior window (away from padding edges) that both share.
let a = yFull[0..., 1_000 ..< 39_000, 0...]
let b = ySub[0..., 1_000 ..< 39_000, 0...]
let maxDiff = abs(a - b).max().item(Float.self)
print("max |yFull - ySub| over shared interior:", maxDiff) // expect ~0 (fp noise)
// Also: the early region of the long output is (wrongly) almost constant.
let early = yFull[0..., 1_000 ..< 30_000, 0...]
print("early-region spread (max-min):", (early.max() - early.min()).item(Float.self))Expected (and observed on macOS GPU / CPU): maxDiff ≈ 0; the early region has normal spread.
Observed on iOS GPU: maxDiff is large; the early region of yFull is nearly constant
(spread ≈ 0), while ySub (same input region, shorter total length) is correct. The boundary
between corrupted and correct output sits at length − 32768.
convTransposed1d shows the same behavior once its output exceeds the threshold.
Workaround
Tile the convolution along the length axis so each call's output stays under 32768 frames, then
concatenate (conv1d is exact when each output tile is fed its full receptive field; transposed conv
via overlap-add). This fully restores correctness on iOS but obviously shouldn't be necessary.
Summary
On the iOS GPU (Apple A-series Metal),
MLX.conv1d(andconvTransposed1d) return incorrectresults when the output length along the convolution axis exceeds 32768 (2¹⁵). Output positions
before roughly
(length − 32768)collapse to a near-constant value; only the last ~32768positions are correct. The same code is correct on the macOS GPU and on CPU.
It looks like a 2¹⁵ addressing/indexing limit in the Metal convolution kernel (block or threadgroup
offset overflowing a 16-bit index).
Impact
Found while running the Kokoro neural TTS vocoder (an iSTFTNet HiFi-GAN-style decoder) on device:
long sentences produced a sustained tone/"beep" because the vocoder's final conv layers emit a
constant spectrogram region. Anything that runs a 1-D conv over a long sequence on iOS (audio
vocoders, long time-series) is affected. The 32768 boundary is well within reach — ~1.4 s of 24 kHz
audio at the post-upsampling resolution.
Environment
mlx)[batch, length, channels],float32.Minimal reproducer (self-contained, no reference output needed)
A 1-D convolution is translation-equivariant, so the first N outputs must not depend on whether
the input is N long or much longer (away from the padding edges). On the iOS GPU they differ:
Expected (and observed on macOS GPU / CPU):
maxDiff ≈ 0; the early region has normal spread.Observed on iOS GPU:
maxDiffis large; the early region ofyFullis nearly constant(spread ≈ 0), while
ySub(same input region, shorter total length) is correct. The boundarybetween corrupted and correct output sits at
length − 32768.convTransposed1dshows the same behavior once its output exceeds the threshold.Workaround
Tile the convolution along the length axis so each call's output stays under 32768 frames, then
concatenate (conv1d is exact when each output tile is fed its full receptive field; transposed conv
via overlap-add). This fully restores correctness on iOS but obviously shouldn't be necessary.