Repository: coreylowman/dfdx (language: rust)

Issue:
Add const generic NUM_THREADS args to launch_cfg
Currently the internal launch_cfg function hardcodes NUM_THREADS to 128

https://github.com/coreylowman/dfdx/blob/main/src/tensor/cuda/mod.rs#L6

This work should:
1. Update the definition to `pub(crate) fn launch_cfg<const NUM_THREADS: u32>(n: u32) -> cudarc::driver::LaunchConfig`
2. Update all usages of launch_cfg to `launch_cfg::<128>`

To verify that all places have been updated (including nightly conv kernels), you can run:
```rust
cargo +nightly check --all -F ci-check,test-cuda
```

If you want to actually run the tests, you can run:
```rust
cargo +nightly test --tests -F test-cuda
```


Discussion / hints:

Fix patch:
diff --git a/src/optim/adam/cuda_kernel.rs b/src/optim/adam/cuda_kernel.rs
index 15d239621..ef1c9b132 100644
--- a/src/optim/adam/cuda_kernel.rs
+++ b/src/optim/adam/cuda_kernel.rs
@@ -68,7 +68,7 @@ where
         let opt_cfg = adam_config_to_cuda(cfg);
         let numel = param.len();
         let func = self.dev.get_func(Self::MOD, Self::FWD).unwrap();
-        let cfg = launch_cfg(numel as u32);
+        let cfg = launch_cfg::<128>(numel as u32);
         let t = <E>::from_i32(t).unwrap();
         let params = (opt_cfg, numel, t, param, moment1, moment2, grad);
         unsafe { func.launch(cfg, params) }?;
diff --git a/src/optim/rmsprop/cuda_kernel.rs b/src/optim/rmsprop/cuda_kernel.rs
index 6d6e18b43..a87430d24 100644
--- a/src/optim/rmsprop/cuda_kernel.rs
+++ b/src/optim/rmsprop/cuda_kernel.rs
@@ -78,7 +78,7 @@ where
         let opt_cfg = rmsprop_config_to_cuda(cfg);
         let numel = param.len();
         let func = self.dev.get_func(Self::MOD, Self::FWD).unwrap();
-        let cfg = launch_cfg(numel as u32);
+        let cfg = launch_cfg::<128>(numel as u32);
         let params = (opt_cfg, numel, param, momentum, square_avg, grad_avg, grad);
         unsafe { func.launch(cfg, params) }?;
         Ok(())
diff --git a/src/optim/sgd/cuda_kernel.rs b/src/optim/sgd/cuda_kernel.rs
index 9216f0870..0e3d37097 100644
--- a/src/optim/sgd/cuda_kernel.rs
+++ b/src/optim/sgd/cuda_kernel.rs
@@ -66,7 +66,7 @@ where
         let opt_cfg = sgd_config_to_cuda(cfg);
         let numel = param.len();
         let func = self.dev.get_func(Self::MOD, Self::FWD).unwrap();
-        let cfg = launch_cfg(numel as u32);
+        let cfg = launch_cfg::<128>(numel as u32);
         unsafe { func.launch(cfg, (opt_cfg, numel, param, velocity, grad)) }?;
         Ok(())
     }
diff --git a/src/tensor/cuda/mod.rs b/src/tensor/cuda/mod.rs
index db328591d..84b7a3395 100644
--- a/src/tensor/cuda/mod.rs
+++ b/src/tensor/cuda/mod.rs
@@ -3,8 +3,7 @@ mod device;
 
 pub use device::{Cuda, CudaError};
 
-pub(crate) fn launch_cfg(n: u32) -> cudarc::driver::LaunchConfig {
-    const NUM_THREADS: u32 = 128;
+pub(crate) fn launch_cfg<const NUM_THREADS: u32>(n: u32) -> cudarc::driver::LaunchConfig {
     let num_blocks = (n + NUM_THREADS - 1) / NUM_THREADS;
     cudarc::driver::LaunchConfig {
         grid_dim: (num_blocks, 1, 1),
diff --git a/src/tensor_ops/attention_reshape/cuda_kernel.rs b/src/tensor_ops/attention_reshape/cuda_kernel.rs
index c0642498c..a051cbd95 100644
--- a/src/tensor_ops/attention_reshape/cuda_kernel.rs
+++ b/src/tensor_ops/attention_reshape/cuda_kernel.rs
@@ -72,7 +72,7 @@ where
             sequence_length,
             past_length,
         };
-        let cfg = launch_cfg(numel as u32);
+        let cfg = launch_cfg::<128>(numel as u32);
         let params = (
             op,
             qkv.data.as_ref(),
diff --git a/src/tensor_ops/axpy/cuda_kernel.rs b/src/tensor_ops/axpy/cuda_kernel.rs
index 6b2e900df..91b371ca8 100644
--- a/src/tensor_ops/axpy/cuda_kernel.rs
+++ b/src/tensor_ops/axpy/cuda_kernel.rs
@@ -33,7 +33,7 @@ where
         }
         let numel = a.len();
         let fwd_fn = self.dev.get_func(Self::FN, Self::FN).unwrap();
-        let cfg = launch_cfg(numel as u32);
+        let cfg = launch_cfg::<128>(numel as u32);
         unsafe { fwd_fn.launch(cfg, (numel, a, alpha, b, beta)) }?;
         Ok(())
     }
diff --git a/src/tensor_ops/boolean/cuda_kernels.rs b/src/tensor_ops/boolean/cuda_kernels.rs
index 32da141da..2f15e3aba 100644
--- a/src/tensor_ops/boolean/cuda_kernels.rs
+++ b/src/tensor_ops/boolean/cuda_kernels.rs
@@ -32,7 +32,7 @@ impl Cuda {
         let rhs_strides: CudaSlice<usize> = self.dev.htod_copy(rhs.strides.into())?;
 
         let fwd_fn = self.dev.get_func(MODULE_NAME, fn_name).unwrap();
-        let cfg = launch_cfg(numel as u32);
+        let cfg = launch_cfg::<128>(numel as u32);
         let params = (
             numel,             // const size_t numel,
             S::NUM_DIMS,       // const size_t num_dims,
@@ -62,7 +62,7 @@ impl BooleanKernel for Cuda {
         let mut storage = unsafe { self.dev.alloc(numel) }?;
 
         let fwd_fn = self.dev.get_func(MODULE_NAME, "boolean_not").unwrap();
-        let cfg = launch_cfg(numel as u32);
+        let cfg = launch_cfg::<128>(numel as u32);
         let params = (
             numel,             // const size_t numel,
             inp.data.as_ref(), // const bool *inp,
diff --git a/src/tensor_ops/choose/cuda_kernel.rs b/src/tensor_ops/choose/cuda_kernel.rs
index d8999bdf6..1fc75bf21 100644
--- a/src/tensor_ops/choose/cuda_kernel.rs
+++ b/src/tensor_ops/choose/cuda_kernel.rs
@@ -47,7 +47,7 @@ where
         let rhs_strides: CudaSlice<usize> = self.dev.htod_copy(rhs.strides.into())?;
 
         let fwd_fn = self.dev.get_func(Self::MOD, Self::FNS[0]).unwrap();
-        let cfg = launch_cfg(numel as u32);
+        let cfg = launch_cfg::<128>(numel as u32);
         let params = (
             numel,              // const size_t numel,
             S::NUM_DIMS,        // const size_t num_dims,
@@ -81,7 +81,7 @@ where
         let cond_strides: CudaSlice<usize> = self.dev.htod_copy(cond.strides.into())?;
         let rhs_strides: CudaSlice<usize> = self.dev.htod_copy(rhs.strides.into())?;
 
-        let cfg = launch_cfg(numel as u32);
+        let cfg = launch_cfg::<128>(numel as u32);
         let params = (
             numel,              // const size_t numel,
             S::NUM_DIMS,        // const size_t num_dims,
diff --git a/src/tensor_ops/cmp/cuda_kernels.rs b/src/tensor_ops/cmp/cuda_kernels.rs
index becf8e739..772cf92d7 100644
--- a/src/tensor_ops/cmp/cuda_kernels.rs
+++ b/src/tensor_ops/cmp/cuda_kernels.rs
@@ -56,7 +56,7 @@ impl<E: Unit, Op: CmpOpCudaKernel<E>> CmpKernel<Op, E> for Cuda {
         let out_strides: CudaSlice<usize> = self.dev.htod_copy(strides.into())?;
 
         let fwd_fn = self.dev.get_func(Op::MODULE_NAME, Op::FWD_FN_NAME).unwrap();
-        let cfg = launch_cfg(numel as u32);
+        let cfg = launch_cfg::<128>(numel as u32);
         let params = (
             numel,             // const size_t numel,
             S::NUM_DIMS,       // const size_t num_dims,
@@ -95,7 +95,7 @@ impl<E: Unit, Op: ScalarCmpOpCudaKernel<E>> ScalarCmpKernel<Op, E> for Cuda {
         let out_strides: CudaSlice<usize> = self.dev.htod_copy(strides.into())?;
 
         let fwd_fn = self.dev.get_func(Op::MODULE_NAME, Op::FWD_FN_NAME).unwrap();
-        let cfg = launch_cfg(numel as u32);
+        let cfg = launch_cfg::<128>(numel as u32);
         let params = (
             numel,             // const size_t numel,
             S::NUM_DIMS,       // const size_t num_dims,
diff --git a/src/tensor_ops/concat/cuda_kernel.rs b/src/tensor_ops/concat/cuda_kernel.rs
index d5a5c8208..f80fc1b60 100644
--- a/src/tensor_ops/concat/cuda_kernel.rs
+++ b/src/tensor_ops/concat/cuda_kernel.rs
@@ -56,14 +56,14 @@ impl<E: Dtype + CudaTypeName> super::ConcatKernel<E> for Cuda {
         {
             let f = self.dev.get_func(&module_name, "concat_bwd").unwrap();
             let numel = grad_a.len();
-            let cfg = launch_cfg(numel as u32);
+            let cfg = launch_cfg::<128>(numel as u32);
             unsafe { f.launch(cfg, (numel, &grad_out.slice(0..numel), grad_a)) }?;
             offset += numel;
         }
         {
             let f = self.dev.get_func(&module_name, "concat_bwd").unwrap();
             let numel = grad_b.len();
-            let cfg = launch_cfg(numel as u32);
+            let cfg = launch_cfg::<128>(numel as u32);
             unsafe { f.launch(cfg, (numel, &grad_out.slice(offset..), grad_b)) }?;
         }
         Ok(())
diff --git a/src/tensor_ops/conv2d/cuda_kernel.rs b/src/tensor_ops/conv2d/cuda_kernel.rs
index 714662628..5933a88ce 100644
--- a/src/tensor_ops/conv2d/cuda_kernel.rs
+++ b/src/tensor_ops/conv2d/cuda_kernel.rs
@@ -80,7 +80,7 @@ where
 
         let img_strides = self.dev.htod_copy(make_4d::<L>(lhs.strides).into())?;
         let unfold_fn = self.dev.get_func(Self::MOD, Self::FNS[0]).unwrap();
-        let cfg = launch_cfg((op.batch * op.chan_in * op.h_out * op.w_out) as u32);
+        let cfg = launch_cfg::<128>((op.batch * op.chan_in * op.h_out * op.w_out) as u32);
         let params = (op, lhs.data.as_ref(), &img_strides, &mut patches);
         unsafe { unfold_fn.launch(cfg, params) }?;
 
@@ -131,7 +131,7 @@ where
         {
             // unfold grad_out into patches
             let unfold_fn = self.dev.get_func(Self::MOD, Self::FNS[1]).unwrap();
-            let cfg = launch_cfg((op.batch * op.chan_out * op.h_in * op.w_in) as u32);
+            let cfg = launch_cfg::<128>((op.batch * op.chan_out * op.h_in * op.w_in) as u32);
             unsafe { unfold_fn.launch(cfg, (op, grad_out, &mut patches)) }?;
         }
 
@@ -139,7 +139,7 @@ where
             // prepare filters for backward operations by
             // swapping dims 0 and 1
             let tr_fn = self.dev.get_func(Self::MOD, Self::FNS[2]).unwrap();
-            let cfg = launch_cfg(rhs.shape.num_elements() as u32);
+            let cfg = launch_cfg::<128>(rhs.shape.num_elements() as u32);
             unsafe {
                 tr_fn.launch_on_stream(
                     self.par_stream.as_ref(),
@@ -195,7 +195,7 @@ where
             // sum all the gradients collected in our broadcasted grad_f
             // into grad_rhs
             let sum_fn = self.dev.get_func(Self::MOD, Self::FNS[3]).unwrap();
-            let cfg = launch_cfg(rhs.shape.num_elements() as u32);
+            let cfg = launch_cfg::<128>(rhs.shape.num_elements() as u32);
             unsafe { sum_fn.launch(cfg, (op, &grad_f_b1023, grad_rhs, &f_strides)) }?;
         }
 
diff --git a/src/tensor_ops/convtrans2d/cuda_kernel.rs b/src/tensor_ops/convtrans2d/cuda_kernel.rs
index 944dc25d4..864abe369 100644
--- a/src/tensor_ops/convtrans2d/cuda_kernel.rs
+++ b/src/tensor_ops/convtrans2d/cuda_kernel.rs
@@ -67,7 +67,7 @@ where
 
         let img_strides = self.dev.htod_copy(make_4d::<L>(lhs.strides).into())?;
         let unfold_fn = self.dev.get_func(Self::MOD, Self::FNS[0]).unwrap();
-        let cfg = launch_cfg((op.batch * op.chan_in * op.h_out * op.w_out) as u32);
+        let cfg = launch_cfg::<128>((op.batch * op.chan_in * op.h_out * op.w_out) as u32);
         let params = (op, lhs.data.as_ref(), &img_strides, &mut patches);
         unsafe { unfold_fn.launch(cfg, params) }?;
 
@@ -117,7 +117,7 @@ where
         {
             // unfold grad_out into patches
             let unfold_fn = self.dev.get_func(Self::MOD, Self::FNS[1]).unwrap();
-            let cfg = launch_cfg((op.batch * op.chan_out * op.h_in * op.w_in) as u32);
+            let cfg = launch_cfg::<128>((op.batch * op.chan_out * op.h_in * op.w_in) as u32);
             unsafe { unfold_fn.launch(cfg, (op, grad_out, &mut patches)) }?;
         }
 
@@ -125,7 +125,7 @@ where
             // prepare filters for backward operations by
             // swapping dims 0 and 1 and adding a batch dimension
             let tr_fn = self.dev.get_func(Self::MOD, Self::FNS[2]).unwrap();
-            let cfg = launch_cfg(rhs.shape.num_elements() as u32);
+            let cfg = launch_cfg::<128>(rhs.shape.num_elements() as u32);
             unsafe {
                 tr_fn.launch_on_stream(
                     self.par_stream.as_ref(),
@@ -181,7 +181,7 @@ where
             // sum all the gradients collected in our broadcasted grad_f
             // into grad_rhs
             let sum_fn = self.dev.get_func(Self::MOD, Self::FNS[3]).unwrap();
-            let cfg = launch_cfg(rhs.shape.num_elements() as u32);
+            let cfg = launch_cfg::<128>(rhs.shape.num_elements() as u32);
             unsafe { sum_fn.launch(cfg, (op, &grad_f_b1023, grad_rhs, &f_strides)) }?;
         }
 
diff --git a/src/tensor_ops/dropout/cuda_kernel.rs b/src/tensor_ops/dropout/cuda_kernel.rs
index b99b35ae7..19c0f8373 100644
--- a/src/tensor_ops/dropout/cuda_kernel.rs
+++ b/src/tensor_ops/dropout/cuda_kernel.rs
@@ -52,7 +52,7 @@ where
         let mut storage = unsafe { self.dev.alloc::<E>(numel) }?;
 
         let fwd_fn = self.dev.get_func(Self::MOD, Self::FNS[0]).unwrap();
-        let cfg = launch_cfg(numel as u32);
+        let cfg = launch_cfg::<128>(numel as u32);
         let params = (op.prob, numel, inp.data.as_ref(), &noise, &mut storage);
         unsafe { fwd_fn.launch(cfg, params) }?;
         Ok(self.build_tensor(inp.shape, inp.strides, storage))
@@ -72,7 +72,7 @@ where
         }?;
         let bwd_fn = self.dev.get_func(Self::MOD, Self::FNS[1]).unwrap();
         let numel = inp.data.len();
-        let cfg = launch_cfg(numel as u32);
+        let cfg = launch_cfg::<128>(numel as u32);
         let params = (op.prob, numel, &noise, grad_inp, grad_out);
         unsafe { bwd_fn.launch(cfg, params) }?;
         Ok(())
diff --git a/src/tensor_ops/max_to/cuda_kernel.rs b/src/tensor_ops/max_to/cuda_kernel.rs
index f7e29f5af..4491f5f0e 100644
--- a/src/tensor_ops/max_to/cuda_kernel.rs
+++ b/src/tensor_ops/max_to/cuda_kernel.rs
@@ -47,7 +47,7 @@ where
         let mut storage = unsafe {
             let mut storage = self.dev.alloc::<E>(dst.num_elements())?;
             fill_fn.launch(
-                launch_cfg(dst.num_elements() as u32),
+                launch_cfg::<128>(dst.num_elements() as u32),
                 (&mut storage, Self::INIT, dst.num_elements()),
             )?;
             storage
@@ -66,7 +66,7 @@ where
             reduction_output_strides::<Ax, Src, Dst>(inp.strides, dst);
         let chunk_len = physical_numel / dst_physical_numel;
 
-        let cfg = launch_cfg(physical_numel as u32);
+        let cfg = launch_cfg::<128>(physical_numel as u32);
 
         let params = (
             physical_numel,    // const size_t numel,
@@ -102,7 +102,7 @@ where
         ))
         .unwrap();
 
-        let cfg = launch_cfg(physical_numel as u32);
+        let cfg = launch_cfg::<128>(physical_numel as u32);
 
         let mut info = Vec::with_capacity(Src::NUM_DIMS * 3);
         info.extend(inp.shape.concrete());
diff --git a/src/tensor_ops/min_to/cuda_kernel.rs b/src/tensor_ops/min_to/cuda_kernel.rs
index 7927056d7..02f09a2c0 100644
--- a/src/tensor_ops/min_to/cuda_kernel.rs
+++ b/src/tensor_ops/min_to/cuda_kernel.rs
@@ -48,7 +48,7 @@ where
         let mut storage = unsafe {
             let mut storage = self.dev.alloc::<E>(dst.num_elements())?;
             fill_fn.launch(
-                launch_cfg(dst.num_elements() as u32),
+                launch_cfg::<128>(dst.num_elements() as u32),
                 (&mut storage, Self::INIT, dst.num_elements()),
             )?;
             storage
@@ -67,7 +67,7 @@ where
             reduction_output_strides::<Ax, Src, Dst>(inp.strides, dst);
         let chunk_len = physical_numel / dst_physical_numel;
 
-        let cfg = launch_cfg(physical_numel as u32);
+        let cfg = launch_cfg::<128>(physical_numel as u32);
         let params = (
             physical_numel,    // const size_t numel,
             num_dims,          // const size_t num_dims,
@@ -103,7 +103,7 @@ where
         ))
         .unwrap();
 
-        let cfg = launch_cfg(physical_numel as u32);
+        let cfg = launch_cfg::<128>(physical_numel as u32);
 
         let mut info = Vec::with_capacity(Src::NUM_DIMS * 3);
         info.extend(inp.shape.concrete());
diff --git a/src/tensor_ops/pool2d/cuda_kernel.rs b/src/tensor_ops/pool2d/cuda_kernel.rs
index 4ede1488a..b38d1f209 100644
--- a/src/tensor_ops/pool2d/cuda_kernel.rs
+++ b/src/tensor_ops/pool2d/cuda_kernel.rs
@@ -35,7 +35,7 @@ macro_rules! pool_impl {
                 let inp_strides = self.dev.htod_copy(make_4d::<I>(inp.strides).into())?;
                 let out_strides = self.dev.htod_copy(make_4d::<O>(out.strides).into())?;
                 let fwd_fn = self.dev.get_func($Fwd, $Fwd).unwrap();
-                let cfg = launch_cfg(out.shape().num_elements() as u32);
+                let cfg = launch_cfg::<128>(out.shape().num_elements() as u32);
                 let params = (
                     op,                           // const Pool2dOp op,
                     &inp_strides,                 // const size_t *inp_strides,
@@ -57,7 +57,7 @@ macro_rules! pool_impl {
                 let inp_strides = self.dev.htod_copy(make_4d::<I>(inp.strides).into())?;
                 let out_strides = self.dev.htod_copy(make_4d::<O>(out.strides).into())?;
                 let bwd_fn = self.dev.get_func($Fwd, $Bwd).unwrap();
-                let cfg = launch_cfg(inp.shape().num_elements() as u32);
+                let cfg = launch_cfg::<128>(inp.shape().num_elements() as u32);
                 let params = (
                     op,                // const Pool2dOp op,
                     &inp_strides,      // const size_t *inp_strides,
diff --git a/src/tensor_ops/reshape_to/cuda_kernel.rs b/src/tensor_ops/reshape_to/cuda_kernel.rs
index edc447c57..0836de83b 100644
--- a/src/tensor_ops/reshape_to/cuda_kernel.rs
+++ b/src/tensor_ops/reshape_to/cuda_kernel.rs
@@ -36,7 +36,7 @@ impl<E: Dtype + CudaTypeName> super::ReshapeKernel<E> for Cuda {
         info.extend(dst.strides());
         let info = self.dev.htod_copy(info)?;
 
-        let cfg = launch_cfg(numel as u32);
+        let cfg = launch_cfg::<128>(numel as u32);
         let params = (
             numel,             // const size_t numel,
             Src::NUM_DIMS,     // const size_t inp_num_dims,
@@ -78,7 +78,7 @@ impl<E: Dtype + CudaTypeName> super::ReshapeKernel<E> for Cuda {
         info.extend(dst.strides());
         let info = self.dev.htod_copy(info)?;
 
-        let cfg = launch_cfg(numel as u32);
+        let cfg = launch_cfg::<128>(numel as u32);
         let params = (
             numel,         // const size_t numel,
             Src::NUM_DIMS, // const size_t inp_num_dims,
diff --git a/src/tensor_ops/roll/cuda_kernel.rs b/src/tensor_ops/roll/cuda_kernel.rs
index 65a6011e5..d237c2da5 100644
--- a/src/tensor_ops/roll/cuda_kernel.rs
+++ b/src/tensor_ops/roll/cuda_kernel.rs
@@ -41,7 +41,7 @@ where
         let out_strides = self.dev.htod_copy(strides.into())?;
 
         let fwd = self.dev.get_func(Self::FNS[0], Self::FNS[0]).unwrap();
-        let cfg = launch_cfg(inp.shape.num_elements() as u32);
+        let cfg = launch_cfg::<128>(inp.shape.num_elements() as u32);
         let params = (
             op,
             S::NUM_DIMS,
@@ -70,7 +70,7 @@ where
         let out_strides = self.dev.htod_copy(strides.into())?;
 
         let bwd = self.dev.get_func(Self::FNS[0], Self::FNS[1]).unwrap();
-        let cfg = launch_cfg(inp.shape.num_elements() as u32);
+        let cfg = launch_cfg::<128>(inp.shape.num_elements() as u32);
         let params = (
             op,
             S::NUM_DIMS,
diff --git a/src/tensor_ops/select_and_gather/cuda_kernel.rs b/src/tensor_ops/select_and_gather/cuda_kernel.rs
index 6f95520d1..913f39cd9 100644
--- a/src/tensor_ops/select_and_gather/cuda_kernel.rs
+++ b/src/tensor_ops/select_and_gather/cuda_kernel.rs
@@ -36,7 +36,7 @@ macro_rules! impl_cuda_kernels {
                 let idx_strides = self.dev.htod_copy(idx.strides.into())?;
 
                 let fwd_fn = self.dev.get_func($GatherMod, $GatherFwd).unwrap();
-                let cfg = launch_cfg(numel as u32);
+                let cfg = launch_cfg::<128>(numel as u32);
                 let params = (
                     numel,             // const size_t numel,
                     inp.data.as_ref(), // const float *inp,
@@ -74,7 +74,7 @@ macro_rules! impl_cuda_kernels {
                 let inp_strides = self.dev.htod_copy(inp.strides.into())?;
                 let idx_strides = self.dev.htod_copy(idx.strides.into())?;
 
-                let cfg = launch_cfg(numel as u32);
+                let cfg = launch_cfg::<128>(numel as u32);
                 let params = (
                     numel,             // const size_t numel,
                     grad_inp,          // float *grad_inp,
@@ -122,7 +122,7 @@ macro_rules! impl_cuda_kernels {
                 let dst_strides = self.dev.htod_copy(dst.strides().into())?;
 
                 let fwd_fn = self.dev.get_func($SelectMod, $SelectFwd).unwrap();
-                let cfg = launch_cfg(numel as u32);
+                let cfg = launch_cfg::<128>(numel as u32);
                 let params = (
                     numel,             // const size_t numel,
                     inp.data.as_ref(), // const float *inp,
@@ -163,7 +163,7 @@ macro_rules! impl_cuda_kernels {
                 let idx_strides = self.dev.htod_copy(idx.strides.into())?;
                 let out_strides = self.dev.htod_copy(out.strides.into())?;
 
-                let cfg = launch_cfg(numel as u32);
+                let cfg = launch_cfg::<128>(numel as u32);
                 let params = (
                     numel,             // const size_t numel,
                     grad_inp,          // float *grad_inp,
diff --git a/src/tensor_ops/slice/cuda_kernel.rs b/src/tensor_ops/slice/cuda_kernel.rs
index 5c328b98d..38f93c29b 100644
--- a/src/tensor_ops/slice/cuda_kernel.rs
+++ b/src/tensor_ops/slice/cuda_kernel.rs
@@ -61,7 +61,7 @@ where
         let strides: CudaSlice<usize> = self.dev.htod_copy(strides.into())?;
 
         let fwd_fn = self.dev.get_func(Self::MOD, Self::FNS[0]).unwrap();
-        let cfg = launch_cfg(numel as u32);
+        let cfg = launch_cfg::<128>(numel as u32);
         let params = (
             numel,             // const size_t numel,
             Src::NUM_DIMS,     // const size_t num_dims,
@@ -97,7 +97,7 @@ where
         let strides: CudaSlice<usize> = self.dev.htod_copy(strides.into())?;
 
         let bwd_fn = self.dev.get_func(Self::MOD, Self::FNS[1]).unwrap();
-        let cfg = launch_cfg(numel as u32);
+        let cfg = launch_cfg::<128>(numel as u32);
         let params = (
             numel,         // const size_t numel,
             Src::NUM_DIMS, // const size_t num_dims,
diff --git a/src/tensor_ops/stack/cuda_kernel.rs b/src/tensor_ops/stack/cuda_kernel.rs
index a28833aa6..3fbbe833f 100644
--- a/src/tensor_ops/stack/cuda_kernel.rs
+++ b/src/tensor_ops/stack/cuda_kernel.rs
@@ -70,7 +70,7 @@ impl<E: Dtype + CudaTypeName> super::StackKernel<E> for Cuda {
         for item in grad_inp.drain(..) {
             let f = self.dev.get_func(&module_name, "stack_bwd").unwrap();
             let numel: usize = item.len();
-            let cfg = launch_cfg(numel as u32);
+            let cfg = launch_cfg::<128>(numel as u32);
             let sub = grad_out.slice(offset..offset + numel);
             unsafe { f.launch(cfg, (numel, &sub, item)) }?;
             offset += numel;
diff --git a/src/tensor_ops/sum_to/cuda_kernel.rs b/src/tensor_ops/sum_to/cuda_kernel.rs
index d9bcbfa83..5d81f5784 100644
--- a/src/tensor_ops/sum_to/cuda_kernel.rs
+++ b/src/tensor_ops/sum_to/cuda_kernel.rs
@@ -61,7 +61,7 @@ where
             reduction_output_strides::<Ax, Src, Dst>(inp.strides, dst);
         let chunk_len = physical_numel / dst_physical_numel;
 
-        let cfg = launch_cfg(physical_numel as u32);
+        let cfg = launch_cfg::<128>(physical_numel as u32);
 
         let mut storage = self.dev.alloc_zeros::<E>(dst.num_elements())?;
         let params = (
@@ -99,7 +99,7 @@ where
         ))
         .unwrap();
 
-        let cfg = launch_cfg(physical_numel as u32);
+        let cfg = launch_cfg::<128>(physical_numel as u32);
 
         let mut info: Vec<usize> = Vec::with_capacity(3 * Src::NUM_DIMS);
         info.extend(inp.shape.concrete());
diff --git a/src/tensor_ops/to_dtype/cuda_kernel.rs b/src/tensor_ops/to_dtype/cuda_kernel.rs
index 3bc01a912..725a56325 100644
--- a/src/tensor_ops/to_dtype/cuda_kernel.rs
+++ b/src/tensor_ops/to_dtype/cuda_kernel.rs
@@ -35,7 +35,12 @@ impl<E1: Unit + CudaTypeName, E2: Unit + CudaTypeName> super::ToDtypeKernel<E1,
 
         let n = inp.data.len();
         let mut out = unsafe { dev.alloc::<E2>(n) }?;
-        unsafe { fwd_fn.launch(launch_cfg(n as u32), (n, inp.data.as_ref(), &mut out)) }?;
+        unsafe {
+            fwd_fn.launch(
+                launch_cfg::<128>(n as u32),
+                (n, inp.data.as_ref(), &mut out),
+            )
+        }?;
 
         Ok(Tensor {
             id: unique_id(),
diff --git a/src/tensor_ops/upscale2d/cuda_kernel.rs b/src/tensor_ops/upscale2d/cuda_kernel.rs
index 2d2013617..2cc7e9049 100644
--- a/src/tensor_ops/upscale2d/cuda_kernel.rs
+++ b/src/tensor_ops/upscale2d/cuda_kernel.rs
@@ -58,7 +58,7 @@ where
 
         let strides = self.dev.htod_copy(make_4d::<I>(inp.strides).into())?;
         let fwd_fn = self.dev.get_func(Self::FWD, Self::FWD).unwrap();
-        let cfg = launch_cfg(out.shape().num_elements() as u32);
+        let cfg = launch_cfg::<128>(out.shape().num_elements() as u32);
         let params = (
             op,
             &strides,
@@ -78,7 +78,7 @@ where
     ) -> Result<(), Self::Err> {
         let strides = self.dev.htod_copy(make_4d::<I>(inp.strides).into())?;
         let bwd_fn = self.dev.get_func(Self::FWD, Self::BWD).unwrap();
-        let cfg = launch_cfg(out.shape().num_elements() as u32);
+        let cfg = launch_cfg::<128>(out.shape().num_elements() as u32);
         let params = (op, &strides, grad_inp, grad_out);
         unsafe { bwd_fn.launch(cfg, params) }?;
         Ok(())
diff --git a/src/tensor_ops/utilities/cuda_kernels.rs b/src/tensor_ops/utilities/cuda_kernels.rs
index 32bf90fc2..f61ba355b 100644
--- a/src/tensor_ops/utilities/cuda_kernels.rs
+++ b/src/tensor_ops/utilities/cuda_kernels.rs
@@ -80,7 +80,7 @@ impl<E: Dtype, K: UnaryOpCudaKernel<E> + DeviceRepr> UnaryKernel<K, E> for Cuda
                 let numel = inp.data.len();
                 let mut storage = unsafe { self.dev.alloc::<E>(numel) }?;
 
-                let cfg = launch_cfg(numel as u32);
+                let cfg = launch_cfg::<128>(numel as u32);
                 let params = (op, numel, inp.data.as_ref(), &mut storage);
                 unsafe { fwd_fn.launch(cfg, params) }?;
 
@@ -96,7 +96,7 @@ impl<E: Dtype, K: UnaryOpCudaKernel<E> + DeviceRepr> UnaryKernel<K, E> for Cuda
             Cow::Owned(mut inp) => {
                 inp.id = unique_id();
                 let numel = inp.data.len();
-                let cfg = launch_cfg(numel as u32);
+                let cfg = launch_cfg::<128>(numel as u32);
                 let params = (op, numel, 0u64, Arc::make_mut(&mut inp.data));
                 unsafe { fwd_fn.launch(cfg, params) }?;
                 Ok(inp)
@@ -115,18 +115,18 @@ impl<E: Dtype, K: UnaryOpCudaKernel<E> + DeviceRepr> UnaryKernel<K, E> for Cuda
         let bwd_fn = self.dev.get_func(K::MODULE_NAME, K::BWD_FN_NAME).unwrap();
         match (inp, out) {
             (Err(inp), Err(_)) => {
-                let cfg = launch_cfg(inp.len as u32);
+                let cfg = launch_cfg::<128>(inp.len as u32);
                 let params = (op, inp.len, 0u64, grad_inp, 0u64, grad_out);
                 unsafe { bwd_fn.launch(cfg, params) }?;
             }
             (Err(inp), Ok(out)) => {
-                let cfg = launch_cfg(inp.len as u32);
+                let cfg = launch_cfg::<128>(inp.len as u32);
                 let params = (op, inp.len, 0u64, grad_inp, out.data.as_ref(), grad_out);
                 unsafe { bwd_fn.launch(cfg, params) }?;
             }
             (Ok(inp), Err(_)) => {
                 let numel = inp.data.len();
-                let cfg = launch_cfg(numel as u32);
+                let cfg = launch_cfg::<128>(numel as u32);
                 let params = (op, numel, inp.data.as_ref(), grad_inp, 0u64, grad_out);
                 unsafe { bwd_fn.launch(cfg, params) }?;
             }
@@ -240,7 +240,7 @@ impl<E: Dtype, K: BinaryOpCudaKernel<E> + DeviceRepr + Clone> BinaryKernel<K, E>
         };
         let strides = shape.strides();
         let numel = shape.num_elements();
-        let cfg = launch_cfg(numel as u32);
+        let cfg = launch_cfg::<128>(numel as u32);
 
         let lhs_strides = match &lhs {
             Cow::Borrowed(lhs) => lhs.strides,
@@ -373,7 +373,7 @@ impl<E: Dtype, K: BinaryOpCudaKernel<E> + DeviceRepr + Clone> BinaryKernel<K, E>
         };
 
         let numel = shape.num_elements();
-        let cfg = launch_cfg(numel as u32);
+        let cfg = launch_cfg::<128>(numel as u32);
 
         let ((out_dims1, out_strides1), rhs_strides1) = permute_for_binary_backward(
             shape.concrete(),


Test patch: