Bound winograd conv2d working set by tiling the batch - #4102
Open
Gusanidas wants to merge 1 commit into
Open
Conversation
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
force-pushed
the
fix/conv2d-large-input
branch
from
August 9, 2026 23:29
52fb201 to
413078a
Compare
Gusanidas
marked this pull request as ready for review
August 10, 2026 00:09
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Addresses the conv2d manifestation of #3979 — the conv-side guard, not a fix for the general silent-error mechanism.
Problem
conv2don 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 exceedsrecommendedMaxWorkingSetSizethe 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 ofrecommendedMaxWorkingSetSizeminus 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_SETinjects the byte budget: covers automatic tiling, both fallbacks, and loud rejection of malformed values.MLX_CONV_WINOGRAD_TILE_BATCHforces 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.