Skip to content

[CUDA] Fix grid overflow in gemm conv unfold kernels for >= 65,536 output positions - #3893

Merged
nastya236 merged 1 commit into
ml-explore:mainfrom
AdamDLuz:fix-gemm-conv-grid-overflow
Jul 24, 2026
Merged

nastya236 merged 1 commit into
ml-explore:mainfrom
AdamDLuz:fix-gemm-conv-grid-overflow

Conversation

@AdamDLuz

@AdamDLuz AdamDLuz commented Jul 22, 2026 •

Copy link
Copy Markdown
Contributor

The naive unfold kernels in backend/cuda/conv/gemm_conv.cu and gemm_grouped_conv.cu launch with

num_blocks.x = cuda::ceil_div(wt_spatial_size, block_dims.x);
num_blocks.y = params.C;
num_blocks.z = mat_M;          // mat_M = N * out_pixels

gridDim.z is limited to 65,535, so any convolution with ≥ 65,536 output positions — e.g. a batch-1 256×256 image, common in VAE decoders — fails with cudaGraphAddKernelNode(...): invalid argument (or cudaLaunchKernelExC(...): invalid argument with MLX_USE_CUDA_GRAPHS=0). Meanwhile grid.x (limit 2³¹−1) carries only ceil_div(wt_spatial_size, block_dims.x), which is 1 for typical filters.

This PR swaps the mapping so the unbounded matrix-row index rides grid.x and the filter-spatial blocks ride grid.z, updating the two tid.x/tid.z consumers in each kernel accordingly. Per-block thread layout and memory access patterns are unchanged.

Same bug class as #3858 (batched matvec, batch > 65,535) and the merged #3666 (binary ops).

Repro (before):

import mlx.core as mx
x = mx.zeros((1, 256, 256, 512), dtype=mx.bfloat16)
w = mx.zeros((512, 3, 3, 512), dtype=mx.bfloat16)
mx.eval(mx.conv2d(x, w, stride=1, padding=1))   # invalid argument
# (1, 128, 128, 512) works: 16,384 <= 65,535

Found while debugging ollama/ollama#14843 (FLUX.2-klein image generation on Linux/CUDA): the VAE decoder's upsample-to-256 conv is the first launch to cross the limit. Verified indirectly on an NVIDIA L4 (CUDA 13) via an equivalent host-side chunking workaround producing correct 512×512 and 1024×1024 images; the patch itself is a pure grid-axis remap (not compile-tested here — missing local dep tree).

Note: num_blocks.y = params.C retains a theoretical >65,535-channel limit, untouched here.

The naive unfold kernels launched with num_blocks.z = mat_M, where
mat_M = N * out_pixels can exceed CUDA's 65,535 grid.y/z limit (a
256x256 output already has 65,536 positions), failing with
cudaGraphAddKernelNode / cuLaunchKernel "invalid argument". Swap the
mapping so the matrix-row index rides grid.x (limit 2^31-1) and the
filter-spatial blocks ride grid.z; per-block thread layout and memory
access patterns are unchanged.

@nastya236 nastya236 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks great! THanks for the fix.

@nastya236
nastya236 merged commit 973e27f into ml-explore:main Jul 24, 2026
28 checks passed
@BrewTestBot BrewTestBot mentioned this pull request Aug 18, 2026
1 task done
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.

2 participants