Fix off_x/off_y typo in steel BaseMMAFrag::load_safe - #3565

Summary

BaseMMAFrag::load_safe in steel/gemm/mma.h uses off_x where off_y belongs in its source-address computation.

On current main:

dst[i * kElemCols + j] =
    static_cast<T>(src[(off_x + i) * str_x + (off_x + j) * str_y]);

The second term should be (off_y + j). This is a typo, not intentional:

• The bounds check on the line directly above uses (off_y + j) < lim_y.
• store_safe and store_slice both compute the address as (off_x + i) * str_x + (off_y + j) * str_y — correct.
• The non-safe load uses src[i * str_x + j * str_y] — correct.

load_safe is the only one of the four that uses off_x twice — and it contradicts its own bounds check, which checks column j against off_y.

Effect

MMATile::load_safe calls MMAFrag::load_safe with off_x = (i * kFragRows) * w_x and off_y = (j * kFragCols) * w_y, which differ for any fragment off the i == j diagonal. The buggy term then folds the row-fragment offset into the column address, so fragments load from the wrong column band — a silent data corruption (no out-of-bounds access).

Concretely, a MMATile<float, 1, 32> loaded via load_safe returns the first column-fragment's two columns replicated across all 32 column-fragments.

I haven't audited MLX's own kernels for exposure — Atile.load_safe call sites (e.g. in quantized.h) route through the device MMATile::load_safe, so flagging for maintainers to check.

Fix

(off_x + j) → (off_y + j).

Credit

The bug was discovered by Claude Opus 4.7 (Anthropic) while building a FlashAttention QK^T GEMM on steel/gemm/mma.h in a downstream project, and traced to this line by static comparison against the three sibling loaders. This PR is opened from a human collaborator's GitHub account as Claude does not have its own.