zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-08-28 00:36:32 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Make threadgroup size less or equal to grid size

Browse Source
This commit is contained in:
Angelos Katharopoulos 2025-08-19 01:13:20 -07:00
parent 432c02dabc
commit 39dbd92df5
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
mlx/backend/cuda/custom_kernel.cpp
View File

@ -352,7 +352,7 @@ void CustomKernel::eval_gpu(
// Make the grid
const auto [tx, ty, tz] = threadgroup_;
const auto [gx, gy, gz] = grid_;
dim3 block(tx, ty, tz);
dim3 block(std::min(tx, gx), std::min(ty, gy), std::min(tz, gz));
dim3 grid((gx + tx - 1) / tx, (gy + ty - 1) / ty, (gz + tz - 1) / tz);
// Call the kernel