Bound winograd conv2d working set by tiling the batch

[Gusanidas]

Addresses the conv2d manifestation of #3979 — the conv-side guard, not a fix for the general silent-error mechanism.

Problem

conv2d on Metal silently returns all zeros for large batched inputs. The Winograd path keeps the padded input and both GEMM workspaces alive at once (~3.7x the input for the reported shapes), all referenced by one command buffer; once that working set exceeds recommendedMaxWorkingSetSize the driver fails the command buffer. As @jasonge27's instrumentation shows, Metal reports the OOM but errors from mid-eval commits are dropped, so nothing surfaces. Winograd's selection criteria (C % 32 == 0 && O % 32 == 0 && C + O >= 256) explain the channel dependence in the report. That dropped-error mechanism is general and needs a separate, complementary fix in the eval machinery; this PR avoids the OOM for conv2d, which also bounds its peak memory.

Fix

Run the batch in tiles sized from the available working set, reusing one set of scratch buffers across tiles — mirroring the row tiling in explicit_gemm_conv_ND_gpu. The budget is 3/4 of recommendedMaxWorkingSetSize minus the conv's own tensors (input, output, and transformed filter), making the tile choice a deterministic function of the conv and the device. Measured failures start at ~99% of the reported limit; the quarter held back is headroom for what is not charged — GPU memory held by other processes and allocations unrelated to the conv.

When the whole batch fits, the common case, the emitted work is unchanged. It falls back to the scratch-free implicit GEMM when not even one batch element fits, or when tiles would carry too few GEMM rows to amortize their fixed cost (small tiles can be far slower than the fallback).

Testing

The real threshold needs tens of GB, so tests shrink the budget instead:

• MLX_CONV_WINOGRAD_WORKING_SET injects the byte budget: covers automatic tiling and both fallbacks.
• MLX_CONV_WINOGRAD_TILE_BATCH forces a tile size (capped by the budget): covers uniform tiles, a short final tile, and a consumer op reading the tiled output in the same eval.

Tiling keeps the per-element reduction order, so tiled Winograd is bit-identical to untiled while the implicit GEMM fallback never is — the tests use exact equality to pin each run to its intended path, with no memory measurements.

At real sizes on a 96 GB M2 Max: the issue's repro passes (previously all-zero at batch 80 on this machine; the reporter hit it at batch 40 on 48 GB), tiled vs untiled outputs are bit-identical at ~50 GB scale, CPU cross-checks agree, and the batch-80 repro peaks at 58 GB where untiled Winograd would need ~84 GB. On the reported fp32 shape (40x512x512, 192 -> 96 channels) Winograd runs ~1.7x faster than implicit GEMM (310 ms vs 536 ms), while a spatially tiny conv (256x4x4, 64 -> 192) cut into one-element Winograd tiles is ~40x slower than implicit GEMM — the two sides of the tile-versus-fallback trade-off. Mid-size Winograd shapes select a single tile and show no regression.

[erwinzhang7]

Verified on an M4 Pro Mac mini with 48 GB, macOS 26.5 — the reporter's
configuration in #3979, which you note you could not test at 96 GB.

Control is 8c28c385, this PR's own parent, since the branch is cut from an
older main. Same sweep on both, 512x512, 192 -> 96, fp32, checked against a CPU
convolution of the same input:

control 8c28c385                      this PR 413078a0
 batch   peak GB   max err  result     batch   peak GB   max err  result
    32      33.7  1.14e-05  ok            32      30.0  1.18e-05  ok
    40      42.1      3.34  ALL ZERO      40      30.1  1.02e-05  ok
    48      50.6      5.64  ALL ZERO      48      29.5  1.69e-05  ok
    64      67.4      5.18  ALL ZERO      64      29.9  2.03e-05  ok

Failure starts at batch 40, the number @kayadibi1 reported on 48 GB. Peak
climbs unbounded on the control, to 67.4 GB asked for on a 48 GB machine. With
the patch it plateaus around 30 GB from batch 32 and stays correct through 64.

Also on an M5 Max via the injected budget: at 512 MB peak drops 1.10 -> 0.36 GB
and the output is bit-identical to untiled, np.array_equal on the uint32
view of 25,165,824 elements. At 1 MB it falls back to implicit GEMM with a max
diff of 0.0398 against untiled winograd, that is the two algorithms differing,
not the tiling. test_conv
passes.

No opinion on the tile-versus-fallback heuristic, which I have not stressed
beyond these shapes.

[zcbenz]

Can you rebase and fix the conflict with https://github.com/ml-explore/mlx/pull/4258/changes#diff-d89046a37f9b2d939417b57314e679e1d9aa2f64a00b0a45b7e4d9656a75056dR903?

[zcbenz]

I think you can just do

Suggested change

	if (auto env_ws = env::get_var("MLX_CONV_WINOGRAD_WORKING_SET", "");
	if (int env_ws = env::get_var("MLX_CONV_WINOGRAD_WORKING_SET", 0);

[zcbenz]

We don't dispatch kernels depending on current memory, it is not reliable and make behavior unpredictable.

[zcbenz]

I would prefer dropping the memory measurements in tests, empirically they tend to be very flaky.

[zcbenz]

Nice work, thanks!