Skip to content

Bound winograd conv2d working set by tiling the batch - #4102

Open
Gusanidas wants to merge 1 commit into
ml-explore:mainfrom
Gusanidas:fix/conv2d-large-input
Open

Bound winograd conv2d working set by tiling the batch#4102
Gusanidas wants to merge 1 commit into
ml-explore:mainfrom
Gusanidas:fix/conv2d-large-input

Conversation

@Gusanidas

@Gusanidas Gusanidas commented Aug 9, 2026

Copy link
Copy Markdown
Contributor

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 MLX's active memory and the transformed filter's allocation: measured failures start at ~99% of the reported limit, and the quarter held back is headroom for memory not charged by this estimate, including pressure from other processes and allocator-cached buffers. Charging all active memory deliberately over-counts unrelated buffers, erring toward smaller tiles or an earlier fallback.

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, both fallbacks, and loud rejection of malformed values.
  • 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.

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.

@Gusanidas
Gusanidas marked this pull request as draft August 9, 2026 20:19
Winograd's scratch is a multiple of the input, all referenced by one
command buffer; once that outgrows the GPU's recommended working set
the kernels silently produce an all-zero output (issue ml-explore#3979). Size a
batch tile from the available working set, reuse one set of scratch
buffers across tiles, and fall back to the scratch-free implicit gemm
when not even one element fits or a tile would carry too little gemm
work to amortize its fixed cost.

The real threshold needs tens of GB, so tests shrink the budget
instead: MLX_CONV_WINOGRAD_WORKING_SET injects the byte budget and
MLX_CONV_WINOGRAD_TILE_BATCH forces a tile size, exercising uniform
and short final tiles, the automatic selector, and both fallbacks at
CI sizes.
@Gusanidas
Gusanidas force-pushed the fix/conv2d-large-input branch from 52fb201 to 413078a Compare August 9, 2026 23:29
@Gusanidas
Gusanidas marked this pull request as ready for review August 10, 2026 00:09
@zcbenz zcbenz added the await verification This pull request is non-trivial and requires a human expert to verify its correctness. label Aug 10, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

await verification This pull request is non-trivial and requires a human expert to verify its correctness.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants