mx.sum over a non-last axis returns all zeros for certain large tensors (e.g. [1, 32, 4096, 131072])

[aidiffuser]

Summary

mx.sum(x, axis=1) returns all zeros instead of the correct sum for x of shape [1, 32, 4096, 131072]. It happens deterministically, for every dtype (float32, float16, bfloat16). The same reduction over the last axis is correct, and smaller shapes are correct.

Repro

import mlx.core as mx

x = mx.ones((1, 32, 4096, 131072), dtype=mx.float32)
y = mx.sum(x, axis=1)
mx.eval(y)
print(float(y.min()), float(y.max()))   # prints 0.0 0.0, expected 32.0 32.0

100% reproducible (5 of 5 fresh runs). Same result with bfloat16 and float16.

What I found

Reducing the last axis is fine: mx.sum(x, axis=3) on the same tensor returns the correct 131072.0. Only the non last axis reduction is affected.

It is size gated. [1, 32, 2048, 131072] through [1, 32, 4032, 131072] are all correct. It first breaks at [1, 32, 4096, 131072] (2^34 elements). Just above that boundary the output is partially zeroed (for example [1, 32, 4352, 131072] gives min 0, max 32), which looks like a 32 bit index overflow in the strided reduction (2^34 is 4 times 2^32).

It is specific to a reduction code path, not a generic element count threshold. At the same 2^34 total, [1, 16, 8192, 131072], [1, 64, 2048, 131072], and [1, 128, 1024, 131072] are all correct. Only the shapes that reduce 32 elements with a large inner size break, e.g. [1, 32, 4096, 131072], [1, 32, 8192, 65536], [1, 32, 16384, 32768].

Expected vs actual

Expected: the correct per position sum (32.0 for the ones tensor above).
Actual: all zeros, or partially zeros just above the boundary.

Workaround

Reducing with elementwise adds (x[:, 0:1] + x[:, 1:2] + ...) or moving the reduced axis to the end both give the correct result.

Environment

MLX 0.32.0.dev20260609+cc3f3e60
Apple M3 Ultra, macOS 26.5.1 (build 25F80)

[aidiffuser]

Real-world motivation: this silently breaks the DeepSeek-V3.2-style DSA "lightning indexer", which reduces over the head axis (axis=1) of a [1, 32, chunk, S] score tensor. With 32 heads and a 4096-token prefill chunk, that tensor hits exactly 2^34 elements at S=131072, so sum(axis=1) returns zeros and the indexer then selects garbage keys. We hit it in GLM-5.2 past ~128k context (worked around by reducing the heads with elementwise adds, as noted above). By the same shape logic, any DSA model with this indexer (for example DeepSeek-V3.2) should be affected at the same boundary.

[varshneydevansh]

Aha I also begun working on it :)

I will drop my changes but this request was
thoroughly detailed thanks.

[aidiffuser]

Since a fix is in progress (thanks @varshneydevansh! 🎉) — here is the repro reshaped as a parametrized regression test, pinning the verified failing shapes and the same-size passing controls so the fix can't regress silently in either direction:

import os
import pytest
import mlx.core as mx

# ~32-36 GiB peak (bf16 input + output); gate for big-memory CI runners only.
BIG = os.environ.get("MLX_RUN_BIG_TESTS") == "1"

FAILING = [  # non-last-axis sum returns (partially) zeroed output
    (1, 32, 4096, 131072),   # canonical: first fully-broken shape (2^34 elems)
    (1, 32, 8192, 65536),    # same pathology, smaller inner dim
    (1, 32, 16384, 32768),
    (1, 32, 4352, 131072),   # just above boundary: PARTIALLY zeroed
]
PASSING_CONTROLS = [  # same 2^34 total elements — must stay correct
    (1, 16, 8192, 131072),
    (1, 64, 2048, 131072),
    (1, 128, 1024, 131072),
    (1, 32, 4032, 131072),   # just below the boundary
]

@pytest.mark.skipif(not BIG, reason="needs ~36 GiB; set MLX_RUN_BIG_TESTS=1")
@pytest.mark.parametrize("shape", FAILING + PASSING_CONTROLS)
def test_non_last_axis_sum_large_strided(shape):
    x = mx.ones(shape, dtype=mx.bfloat16)  # bf16 = lowest footprint; fp16/fp32 equally affected
    y = mx.sum(x, axis=1)
    mx.eval(y)
    expected = float(shape[1])
    assert float(y.min()) == expected and float(y.max()) == expected, (
        f"sum(axis=1) over {shape}: min={float(y.min())}, max={float(y.max())}, "
        f"expected uniform {expected}"
    )

Notes for whoever lands the fix:

• The (1, 32, 4352, 131072) case matters: just above the 2^34 boundary the output is only partially zeroed, so a min()==max()==expected assertion (not just "not all zeros") is the right strictness.
• All three of bf16/fp16/fp32 reproduce; bf16 keeps the test's footprint lowest (~32 GiB input). If that is still too heavy for CI, the three FAILING shapes are interchangeable evidence-wise — keeping only (1, 32, 16384, 32768) plus one control still pins the bug, though the exact 2^34 boundary pair (4032/4096 × 131072) is the most diagnostic if the root cause is the suspected 32-bit index overflow in the strided-reduction path.
• Verified on 0.32.0.dev20260609+cc3f3e60 / M3 Ultra; happy to re-run any candidate fix against the full shape matrix on 512 GB hardware — including the downstream 128k+ end-to-end case (GLM-5.2 DSA indexer) that this bug silently corrupted.

🤖 Generated with Claude Code