refactor: Move Runtime to cubecl-runtime (#1068)

Author: wingertge

Cargo.toml

@@ -103,7 +103,7 @@ tracel-llvm = { version = "20.1.4-5", features = ["mlir-helpers"] }
 # tracel-llvm = { git = "https://github.com/tracel-ai/tracel-llvm.git", branch = "fix/linux", package = "tracel-llvm", features = ["mlir-helpers"] }
 # tracel-llvm = { path = "../tracel-llvm/crates/tracel-llvm", features = ["mlir-helpers"] }
 
-cudarc = { version = "0.17.7", features = [
+cudarc = { version = "0.18.1", features = [
     "std",
     "driver",
     "nvrtc",

crates/cubecl-attention/src/base.rs

@@ -25,15 +25,15 @@ pub enum Strategy {
 #[allow(clippy::result_large_err, clippy::too_many_arguments)]
 pub fn launch<R: Runtime>(
     strategy: &Strategy,
-    client: &ComputeClient<R::Server>,
+    client: &ComputeClient<R>,
     query: TensorHandle<R>,
     key: TensorHandle<R>,
     value: TensorHandle<R>,
     mask: Option<TensorHandle<R>>,
     out: TensorHandle<R>,
     attention_elems: AttentionElems,
 ) -> Result<(), AttentionSetupError> {
-    launch_ref::<R>(
+    launch_ref(
         strategy,
         client,
         &query.as_ref(),
@@ -48,7 +48,7 @@ pub fn launch<R: Runtime>(
 #[allow(clippy::result_large_err, clippy::too_many_arguments)]
 pub fn launch_ref<R: Runtime>(
     strategy: &Strategy,
-    client: &ComputeClient<R::Server>,
+    client: &ComputeClient<R>,
     query: &TensorHandleRef<R>,
     key: &TensorHandleRef<R>,
     value: &TensorHandleRef<R>,
@@ -79,15 +79,16 @@ pub fn launch_ref<R: Runtime>(
 }
 
 pub fn launch_attention<R: Runtime, A: Algorithm>(
-    client: &ComputeClient<R::Server>,
+    client: &ComputeClient<R>,
     query: &TensorHandleRef<R>,
     key: &TensorHandleRef<R>,
     value: &TensorHandleRef<R>,
     mask: &Option<TensorHandleRef<R>>,
     out: &TensorHandleRef<R>,
     attention_elems: &AttentionElems,
 ) -> Result<(), AttentionSetupError> {
-    let line_sizes = AvailableLineSizes::from_elem_types::<R>(
+    let line_sizes = AvailableLineSizes::from_elem_types(
+        client,
         query.elem_size,
         attention_elems.mask.size(),
         out.elem_size,
@@ -135,7 +136,7 @@ pub fn launch_attention<R: Runtime, A: Algorithm>(
         two_rows_in_array_tile: false,
     };
 
-    let config = BlackboxAcceleratedAlgorithm::setup::<R>(
+    let config = BlackboxAcceleratedAlgorithm::setup(
         client,
         &problem,
         &selection,

crates/cubecl-attention/src/components/batch/base.rs

@@ -27,7 +27,7 @@ pub trait BatchAttentionFamily: Send + Sync + 'static {
     /// Out-of-bounds can happen
     #[allow(clippy::too_many_arguments)]
     unsafe fn launch_unchecked<'a, AA: AttentionArgs, R: Runtime>(
-        client: &ComputeClient<<R as Runtime>::Server>,
+        client: &ComputeClient<R>,
         cube_dim: CubeDim,
         cube_count: CubeCount,
         input: InputRuntimeArg<'a, AA, R>,
@@ -41,7 +41,7 @@ pub trait BatchAttentionFamily: Send + Sync + 'static {
     ///
     /// This function may return an error if the configuration cannot be supported on the current runtime.
     fn setup<R: Runtime>(
-        client: &ComputeClient<R::Server>,
+        client: &ComputeClient<R>,
         problem: &AttentionProblem,
         selection: &AttentionSelection,
         line_sizes: &AttentionLineSizes,

crates/cubecl-attention/src/components/batch/simple/setup.rs

@@ -23,13 +23,13 @@ impl<GA: GlobalAttentionFamily> BatchAttentionFamily for SimpleBatchAttentionFam
     type Config = SimpleBatchConfig<GA::Config>;
 
     fn setup<R: cubecl_core::Runtime>(
-        client: &ComputeClient<R::Server>,
+        client: &ComputeClient<R>,
         problem: &AttentionProblem,
         selection: &AttentionSelection,
         line_sizes: &AttentionLineSizes,
         dtypes: &AttentionElems,
     ) -> Result<Self::Config, crate::components::AttentionSetupError> {
-        let global_config = GA::setup::<R>(client, problem, selection, line_sizes, dtypes)?;
+        let global_config = GA::setup(client, problem, selection, line_sizes, dtypes)?;
 
         SimpleBatchConfig::new(
             global_config,
@@ -41,7 +41,7 @@ impl<GA: GlobalAttentionFamily> BatchAttentionFamily for SimpleBatchAttentionFam
     }
 
     unsafe fn launch_unchecked<'a, AA: AttentionArgs, R: cubecl_core::Runtime>(
-        client: &cubecl_core::prelude::ComputeClient<<R as cubecl_core::Runtime>::Server>,
+        client: &cubecl_core::prelude::ComputeClient<R>,
         cube_dim: cubecl_core::CubeDim,
         cube_count: cubecl_core::CubeCount,
         input: InputRuntimeArg<'a, AA, R>,

crates/cubecl-attention/src/components/global/base.rs

@@ -23,7 +23,7 @@ pub trait GlobalAttentionFamily: Send + Sync + 'static {
     ///
     /// This function may return an error if the configuration cannot be supported on the current runtime.
     fn setup<R: Runtime>(
-        client: &ComputeClient<R::Server>,
+        client: &ComputeClient<R>,
         problem: &AttentionProblem,
         selection: &AttentionSelection,
         line_sizes: &AttentionLineSizes,

crates/cubecl-attention/src/components/global/simple/setup.rs

@@ -40,13 +40,13 @@ impl<
     type Config = SimpleGlobalAttentionConfig<SA::Config>;
 
     fn setup<R: cubecl_core::Runtime>(
-        client: &ComputeClient<R::Server>,
+        client: &ComputeClient<R>,
         problem: &AttentionProblem,
         selection: &AttentionSelection,
         line_sizes: &AttentionLineSizes,
         dtypes: &AttentionElems,
     ) -> Result<Self::Config, AttentionSetupError> {
-        let stage_config = SA::setup::<R>(client, problem, selection, line_sizes, dtypes)?;
+        let stage_config = SA::setup(client, problem, selection, line_sizes, dtypes)?;
 
         let precompute_job = LoadingPrecomputeStrategy::Never.into();
         let plane_dim = stage_config.plane_dim();

crates/cubecl-attention/src/components/line_size.rs

@@ -1,6 +1,6 @@
 use std::fmt::Debug;
 
-use cubecl_core::{LineSizeError, Runtime, tensor_line_size_parallel};
+use cubecl_core::{LineSizeError, Runtime, client::ComputeClient, tensor_line_size_parallel};
 
 use crate::components::{AttentionIdent, AttentionSetupError};
 
@@ -29,10 +29,17 @@ pub struct AvailableLineSizes {
 }
 
 impl AvailableLineSizes {
-    pub fn from_elem_types<R: Runtime>(elem_in: usize, elem_mask: usize, elem_out: usize) -> Self {
-        let in_available: Vec<u8> = R::io_optimized_line_sizes_unchecked(elem_in).collect();
-        let mask_available: Vec<u8> = R::io_optimized_line_sizes_unchecked(elem_mask).collect();
-        let out_available = R::io_optimized_line_sizes_unchecked(elem_out).collect();
+    pub fn from_elem_types<R: Runtime>(
+        client: &ComputeClient<R>,
+        elem_in: usize,
+        elem_mask: usize,
+        elem_out: usize,
+    ) -> Self {
+        let in_available: Vec<u8> = client.io_optimized_line_sizes_unchecked(elem_in).collect();
+        let mask_available: Vec<u8> = client
+            .io_optimized_line_sizes_unchecked(elem_mask)
+            .collect();
+        let out_available = client.io_optimized_line_sizes_unchecked(elem_out).collect();
 
         AvailableLineSizes {
             query: in_available.clone(),

crates/cubecl-attention/src/components/stage/base.rs

@@ -45,7 +45,7 @@ pub trait StageAttentionFamily: Send + Sync + 'static {
     ///
     /// This function may return an error if the configuration cannot be supported on the current runtime.
     fn setup<R: Runtime>(
-        client: &ComputeClient<R::Server>,
+        client: &ComputeClient<R>,
         problem: &AttentionProblem,
         selection: &AttentionSelection,
         line_sizes: &AttentionLineSizes,

crates/cubecl-attention/src/components/stage/plane/setup.rs

@@ -53,7 +53,7 @@ impl<
     type Config = PartitionAttentionConfig<TA::Config>;
 
     fn setup<R: cubecl_core::Runtime>(
-        client: &ComputeClient<R::Server>,
+        client: &ComputeClient<R>,
         problem: &AttentionProblem,
         selection: &AttentionSelection,
         line_sizes: &AttentionLineSizes,
@@ -62,8 +62,7 @@ impl<
         let num_planes = selection.tiling_scheme.stage_size.seq_q
             * TA::computation_resources()?.num_planes(selection.plane_dim)?;
 
-        let tile_config =
-            TA::setup::<R>(client, problem, selection, line_sizes, num_planes, dtypes)?;
+        let tile_config = TA::setup(client, problem, selection, line_sizes, num_planes, dtypes)?;
 
         let key_smem_config = StageMemoryConfig {
             num_planes,

crates/cubecl-attention/src/components/stage/unit/setup.rs

@@ -53,7 +53,7 @@ impl<
     type Config = PartitionAttentionConfig<TA::Config>;
 
     fn setup<R: cubecl_core::Runtime>(
-        client: &ComputeClient<R::Server>,
+        client: &ComputeClient<R>,
         problem: &AttentionProblem,
         selection: &AttentionSelection,
         line_sizes: &AttentionLineSizes,
@@ -70,8 +70,7 @@ impl<
         };
 
         let num_planes = compute_resources.num_planes(selection.plane_dim)?;
-        let tile_config =
-            TA::setup::<R>(client, problem, selection, line_sizes, num_planes, dtypes)?;
+        let tile_config = TA::setup(client, problem, selection, line_sizes, num_planes, dtypes)?;
 
         let key_smem_config = StageMemoryConfig {
             num_planes,