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

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)

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.

Aha I also begun working on it :)

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

@varshneydevansh

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

fix(deepseek_v32): work around MLX non-last-axis sum bug corrupting D…

I think the root cause here is a Metal dispatch grid-dimension overflow on the host side, not the reduction kernel math itself.

In mlx/backend/metal/reduce.cpp, strided_reduce_looped builds the grid roughly like this (around L763):

MTL::Size grid_dims(
    threadgroup_size * ((args.reduction_stride + BN - 1) / BN),
    out_grid_size.width, out_grid_size.height);

so the x dimension is threadgroup_size * ceil(reduction_stride / BN) (= 256 * ceil(reduction_stride / 32)), passed straight to dispatch. This path does not go through get_2d_grid_dims_common (mlx/backend/common/utils.cpp), which is what caps each grid dimension below UINT32_MAX elsewhere and factors an oversized extent into 2D.

For the reported case (reducing a non-last axis of [1, 32, 4096, 131072]), reduction_stride is the trailing 4096 * 131072 = 2^29, so grid-x works out to 256 * (2^29 / 32) = 2^32 - right at Metal's per-dimension dispatch limit. Statically walking the 8 shapes in the thread lines up with the report: the failing ones evaluate to 2^32 (and the partially-zeroed case to 2^32 + 2^28, which wraps mod 2^32 to 2^28, matching the partial corruption), while the passing ones stay below 2^32 or fall into strided_reduce_small instead.

I have not been able to confirm this at runtime (no Metal GPU on hand), so this is a static read of the dispatch arithmetic rather than a verified fix. If it is correct, the fix would be to route this dispatch's grid through the same 2D/3D factorization used elsewhere (with matching gid reconstruction in col_reduce_looped / col_reduce_2pass) instead of emitting a single oversized x dimension. Happy to help if a maintainer can confirm the direction.

glm5_next: gathered sparse prefill + fp32 matmul indexer contraction

Navigation Menu

Repository navigation

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

Summary

Repro

What I found

Expected vs actual

Workaround

Environment

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

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

Description

Summary

Repro

What I found

Expected vs actual

Workaround

Environment

Activity

aidiffuser commented on Jun 30, 2026

varshneydevansh commented on Jul 2, 2026

aidiffuser commented on Jul 2, 2026

devYRPauli commented on Jul 16, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions