Skip to content

Fix conv_transpose maxBufferLength failures on Metal via tiled unfold#3845

Open
eyupcanakman wants to merge 1 commit into
ml-explore:mainfrom
eyupcanakman:fix/conv-transpose-tiled-unfold
Open

Fix conv_transpose maxBufferLength failures on Metal via tiled unfold#3845
eyupcanakman wants to merge 1 commit into
ml-explore:mainfrom
eyupcanakman:fix/conv-transpose-tiled-unfold

Conversation

@eyupcanakman

Copy link
Copy Markdown
Contributor

Proposed changes

conv_transpose1d/2d/3d (and conv_general with input_dilation > 1) on Metal go through the explicit GEMM path, which allocates one full unfold buffer of implicit_M by implicit_K. For large outputs that buffer exceeds the device maxBufferLength and the allocation throws, even when the output itself fits in memory (#3082).

The explicit GEMM now unfolds and runs in row tiles that each stay under maxBufferLength, reusing a single unfold buffer so peak memory is one tile. Both the grouped and non-grouped paths are covered. The approach follows #3084.

MLX_CONV_UNFOLD_TILE_ROWS caps the tile size so the tiled path can be exercised on small inputs, which the added test uses.

Checklist

Put an x in the boxes that apply.

  • I have read the CONTRIBUTING document
  • I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (if needed)

conv_transpose on Metal allocated one full unfold buffer that could exceed the device maxBufferLength for large outputs. Unfold and run the GEMM in row tiles reusing a single buffer, so peak memory stays at one tile (ml-explore#3082).
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant