zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Check for valid launch parameters

Browse Source
This commit is contained in:
Angelos Katharopoulos
2025-06-06 11:41:56 -07:00
parent 570dd8287a
commit 97bd67c032
1 changed files with 14 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
mlx/backend/metal/normalization.cpp
View File

@@ -273,7 +273,13 @@ void LayerNorm::eval_gpu(
size_t threadgroup_needed = (axis_size + n_reads - 1) / n_reads;
size_t simds_needed = (threadgroup_needed + simd_size - 1) / simd_size;
size_t threadgroup_size = simd_size * simds_needed;
assert(threadgroup_size <= kernel->maxTotalThreadsPerThreadgroup());
if (threadgroup_size > kernel->maxTotalThreadsPerThreadgroup()) {
std::ostringstream msg;
msg << "[layer_norm] Threadgroup size " << threadgroup_size
<< " is larger than the maximum allowed threadgroup size "
<< kernel->maxTotalThreadsPerThreadgroup();
throw std::runtime_error(msg.str());
}
size_t n_threads = n_rows * threadgroup_size;
grid_dims = MTL::Size(n_threads, 1, 1);
group_dims = MTL::Size(threadgroup_size, 1, 1);
@@ -396,7 +402,13 @@ void LayerNormVJP::eval_gpu(
size_t threadgroup_needed = (axis_size + n_reads - 1) / n_reads;
size_t simds_needed = (threadgroup_needed + simd_size - 1) / simd_size;
size_t threadgroup_size = simd_size * simds_needed;
assert(threadgroup_size <= kernel->maxTotalThreadsPerThreadgroup());
if (threadgroup_size > kernel->maxTotalThreadsPerThreadgroup()) {
std::ostringstream msg;
msg << "[vjp_layer_norm] Threadgroup size " << threadgroup_size
<< " is larger than the maximum allowed threadgroup size "
<< kernel->maxTotalThreadsPerThreadgroup();
throw std::runtime_error(msg.str());
}
size_t n_threads = n_rows * threadgroup_size;
grid_dims = MTL::Size(n_threads, 1, 1);
group_dims = MTL::Size(threadgroup_size, 1, 1);