diff --git a/ext/AMDGPUExt/AMDGPUExt.jl b/ext/AMDGPUExt/AMDGPUExt.jl
index 0524ade..047eef8 100644
--- a/ext/AMDGPUExt/AMDGPUExt.jl
+++ b/ext/AMDGPUExt/AMDGPUExt.jl
@@ -105,7 +105,7 @@ function _parallel_for(indexer::TI, f, (m, n), (M, N), x...) where {TI}
     kernel, shmem_size = _kernel_maxshmem(_parallel_for_amdgpu_MN, kargs)
     config = AMDGPU.launch_configuration(kernel; shmem = shmem_size)
     maxThreadsX = sqrt(config.groupsize)
-    y_thr = floor(Int, (n / m) * maxThreadsX)
+    y_thr = clamp(floor(Int, (n / m) * maxThreadsX), 1, config.groupsize)
     x_thr = fld(config.groupsize, y_thr)
     threads = (x_thr, y_thr)
     blocks = (cld(m, x_thr), cld(n, y_thr))
@@ -118,7 +118,7 @@ function JACC.parallel_for(
     dev = AMDGPU.device()
     props = AMDGPU.HIP.properties(dev)
     maxBlocks = (x = props.maxGridSize[1], y = props.maxGridSize[2])
-    if M < N && maxBlocks.x > maxBlocks.y
+    if M < N && maxBlocks.x >= maxBlocks.y
         _parallel_for(BlockIndexerSwapped(), f, (N, M), (M, N), x...)
     else
         _parallel_for(BlockIndexerBasic(), f, (M, N), (M, N), x...)
@@ -137,7 +137,7 @@ function _parallel_for(indexer::TI, f, spec::LaunchSpec{AMDGPUBackend}, (m, n),
     if spec.threads == 0
         config = AMDGPU.launch_configuration(kernel; shmem = spec.shmem_size)
         maxThreadsX = sqrt(config.groupsize)
-        y_thr = floor(Int, (n / m) * maxThreadsX)
+        y_thr = clamp(floor(Int, (n / m) * maxThreadsX), 1, config.groupsize)
         x_thr = fld(config.groupsize, y_thr)
         spec.threads = (x_thr, y_thr)
     end
@@ -158,7 +158,7 @@ function JACC.parallel_for(
     dev = AMDGPU.device()
     props = AMDGPU.HIP.properties(dev)
     maxBlocks = (x = props.maxGridSize[1], y = props.maxGridSize[2])
-    if M < N && maxBlocks.x > maxBlocks.y
+    if M < N && maxBlocks.x >= maxBlocks.y
         _parallel_for(BlockIndexerSwapped(), f, spec, (N, M), (M, N), x...)
     else
         _parallel_for(BlockIndexerBasic(), f, spec, (M, N), (M, N), x...)
diff --git a/ext/CUDAExt/CUDAExt.jl b/ext/CUDAExt/CUDAExt.jl
index 74e80ae..27caa1a 100644
--- a/ext/CUDAExt/CUDAExt.jl
+++ b/ext/CUDAExt/CUDAExt.jl
@@ -100,7 +100,7 @@ function _parallel_for(indexer::TI, f, (m, n), (M, N), x...) where {TI}
     kernel, shmem_size = _kernel_maxshmem(_parallel_for_cuda_MN, kargs)
     config = CUDA.launch_configuration(kernel.fun; shmem = shmem_size)
     maxThreadsX = sqrt(config.threads)
-    y_thr = floor(Int, (n / m) * maxThreadsX)
+    y_thr = clamp(floor(Int, (n / m) * maxThreadsX), 1, config.threads)
     x_thr = fld(config.threads, y_thr)
     threads = (x_thr, y_thr)
     blocks = (cld(m, x_thr), cld(n, y_thr))
@@ -114,7 +114,7 @@ function JACC.parallel_for(f, ::CUDABackend, (M, N)::NTuple{2, Integer}, x...)
         x = attribute(dev, CUDA.DEVICE_ATTRIBUTE_MAX_GRID_DIM_X),
         y = attribute(dev, CUDA.DEVICE_ATTRIBUTE_MAX_GRID_DIM_Y)
     )
-    if M < N && maxBlocks.x > maxBlocks.y
+    if M < N && maxBlocks.x >= maxBlocks.y
         _parallel_for(BlockIndexerSwapped(), f, (N, M), (M, N), x...)
     else
         _parallel_for(BlockIndexerBasic(), f, (M, N), (M, N), x...)
@@ -133,7 +133,7 @@ function _parallel_for(indexer::TI, f, spec::LaunchSpec{CUDABackend}, (m, n),
     if spec.threads == 0
         config = CUDA.launch_configuration(kernel.fun; shmem = spec.shmem_size)
         maxThreadsX = sqrt(config.threads)
-        y_thr = floor(Int, (n / m) * maxThreadsX)
+        y_thr = clamp(floor(Int, (n / m) * maxThreadsX), 1, config.threads)
         x_thr = fld(config.threads, y_thr)
         spec.threads = (x_thr, y_thr)
     end
@@ -156,7 +156,7 @@ function JACC.parallel_for(
         x = attribute(dev, CUDA.DEVICE_ATTRIBUTE_MAX_GRID_DIM_X),
         y = attribute(dev, CUDA.DEVICE_ATTRIBUTE_MAX_GRID_DIM_Y)
     )
-    if M < N && maxBlocks.x > maxBlocks.y
+    if M < N && maxBlocks.x >= maxBlocks.y
         _parallel_for(BlockIndexerSwapped(), f, spec, (N, M), (M, N), x...)
     else
         _parallel_for(BlockIndexerBasic(), f, spec, (M, N), (M, N), x...)
diff --git a/ext/oneAPIExt/oneAPIExt.jl b/ext/oneAPIExt/oneAPIExt.jl
index cdec47d..a4cce09 100644
--- a/ext/oneAPIExt/oneAPIExt.jl
+++ b/ext/oneAPIExt/oneAPIExt.jl
@@ -71,7 +71,7 @@ function _parallel_for(indexer::TI, f, (m, n), (M, N), x...) where {TI}
     kernel = @oneapi launch=false _parallel_for_oneapi_MN(indexer, (M, N), f, x...)
     maxThreads = div(oneAPI.launch_configuration(kernel), 2)
     maxThreadsX = sqrt(maxThreads)
-    y_thr = floor(Int, (n / m) * maxThreadsX)
+    y_thr = clamp(floor(Int, (n / m) * maxThreadsX), 1, maxThreads)
     x_thr = fld(maxThreads, y_thr)
     items = (x_thr, y_thr)
     groups = (cld(m, items[1]), cld(n, items[2]))
@@ -85,7 +85,7 @@ function JACC.parallel_for(
     dev = oneAPI.device()
     props = oneAPI.compute_properties(dev)
     maxBlocks = (x = props.maxGroupCountX, y = props.maxGroupCountY)
-    if M < N && maxBlocks.x > maxBlocks.y
+    if M < N && maxBlocks.x >= maxBlocks.y
         _parallel_for(BlockIndexerSwapped(), f, (N, M), (M, N), x...)
     else
         _parallel_for(BlockIndexerBasic(), f, (M, N), (M, N), x...)
@@ -99,7 +99,7 @@ function _parallel_for(indexer::TI, f, spec::LaunchSpec{oneAPIBackend}, (m, n),
     if spec.threads == 0
         maxThreads = oneAPI.launch_configuration(kernel)
         maxThreadsX = sqrt(maxThreads)
-        y_thr = floor(Int, (n / m) * maxThreadsX)
+        y_thr = clamp(floor(Int, (n / m) * maxThreadsX), 1, maxThreads)
         x_thr = fld(maxThreads, y_thr)
         spec.threads = (x_thr, y_thr)
     end
@@ -120,7 +120,7 @@ function JACC.parallel_for(
     dev = oneAPI.device()
     props = oneAPI.compute_properties(dev)
     maxBlocks = (x = props.maxGroupCountX, y = props.maxGroupCountY)
-    if M < N && maxBlocks.x > maxBlocks.y
+    if M < N && maxBlocks.x >= maxBlocks.y
         _parallel_for(BlockIndexerSwapped(), f, spec, (N, M), (M, N), x...)
     else
         _parallel_for(BlockIndexerBasic(), f, spec, (M, N), (M, N), x...)
diff --git a/test/unittests.jl b/test/unittests.jl
index 371bfa4..f51a221 100644
--- a/test/unittests.jl
+++ b/test/unittests.jl
@@ -571,6 +571,27 @@ end
     @test JACC.to_host(C)≈C_expected rtol=1e-5
 end
 
+@testset "Add-2D imbalanced" begin
+    function add!(i, j, A, B, C)
+        @inbounds C[i, j] = A[i, j] + B[i, j]
+    end
+
+    for (M, N) in ((1024, 16), (16, 1024))
+        C_expected = Float32(2.0) .* ones(Float32, M, N)
+
+        A = JACC.ones(Float32, M, N)
+        B = JACC.ones(Float32, M, N)
+
+        C = JACC.zeros(Float32, M, N)
+        JACC.parallel_for((M, N), add!, A, B, C)
+        @test JACC.to_host(C)≈C_expected rtol=1e-5
+
+        C_spec = JACC.zeros(Float32, M, N)
+        JACC.parallel_for(JACC.launch_spec(), (M, N), add!, A, B, C_spec)
+        @test JACC.to_host(C_spec)≈C_expected rtol=1e-5
+    end
+end
+
 @testset "Add-3D" begin
     function add!(i, j, k, A, B, C)
         @inbounds C[i, j, k] = A[i, j, k] + B[i, j, k]