Chore/update/cubecl (#2067)

Author: nathanielsimard

Cargo.lock

@@ -16,7 +16,7 @@ members = [
 
 exclude = [
     "examples/notebook",
-    "crates/burn-cuda",     # comment this line to work on burn-cuda
+    # "crates/burn-cuda",     # comment this line to work on burn-cuda
 ]
 
 [workspace.package]
@@ -93,7 +93,7 @@ crossterm = "0.27.0"
 
 # WGPU stuff
 text_placeholder = "0.5.1"
-wgpu = "0.20.1"
+wgpu = "22.0.0"
 
 # Benchmarks and Burnbench
 arboard = "3.4.0"
@@ -140,8 +140,12 @@ nvml-wrapper = "0.10.0"
 sysinfo = "0.30.13"
 systemstat = "0.2.3"
 
+### For the main burn branch. ###
 cubecl = { version = "0.1.1", git = "https://github.com/tracel-ai/cubecl", default-features = false }
 cubecl-common = { version = "0.1.1", git = "https://github.com/tracel-ai/cubecl", default-features = false }
+### For local development. ###
+# cubecl = { path = "../cubecl/crates/cubecl" }
+# cubecl-common = { path = "../cubecl/crates/cubecl-common" }
 
 [profile.dev]
 debug = 0 # Speed up compilation time and not necessary.

Cargo.toml

@@ -32,7 +32,7 @@ tokio = { workspace = true, optional = true }
 
 # Parallel
 rayon = { workspace = true, optional = true }
-cubecl-common = { workspace = true, default-features = false }
+cubecl-common = { workspace = true }
 
 [dev-dependencies]
 dashmap = { workspace = true }

crates/burn-common/Cargo.toml

@@ -124,8 +124,8 @@ impl<B: Backend> RotaryEncoding<B> {
     ///
     /// Arguments:
     /// * `x` - Input tensor of shape (..., seq_len, d_model). Accommodate both 3D and 4D tensors
-    /// for (batch size, seq_len, hidden_dim) or (batch size, num_heads, seq_len, hidden_dim)
-    /// respectively.
+    ///    for (batch size, seq_len, hidden_dim) or (batch size, num_heads, seq_len, hidden_dim)
+    ///    respectively.
     ///
     /// Returns:
     /// * Output tensor with the same shape as input tensor after applying rotary encoding.
@@ -139,8 +139,8 @@ impl<B: Backend> RotaryEncoding<B> {
     ///
     /// Arguments:
     /// * `x` - Input tensor of shape (..., seq_len, d_model). Accommodate both 3D and 4D tensors
-    /// for (batch size, seq_len, hidden_dim) or (batch size, num_heads, seq_len, hidden_dim)
-    /// respectively.
+    ///    for (batch size, seq_len, hidden_dim) or (batch size, num_heads, seq_len, hidden_dim)
+    ///    respectively.
     /// * `start` - Sequence start position index.
     ///
     /// Returns:

crates/burn-core/src/nn/rope_encoding.rs

@@ -11,11 +11,11 @@ repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-cuda"
 version.workspace = true
 
 [features]
-default = ["fusion", "burn-jit/default"]
+default = ["fusion", "burn-jit/default", "cubecl/default"]
 fusion = ["burn-fusion", "burn-jit/fusion"]
 autotune = ["burn-jit/autotune"]
 doc = ["burn-jit/doc"]
-std = ["burn-jit/std"]
+std = ["burn-jit/std", "cubecl/std"]
 
 [dependencies]
 cubecl = { workspace = true, features = ["cuda"] }

crates/burn-cuda/Cargo.toml

@@ -73,17 +73,17 @@ impl From<&'static str> for SqliteDatasetError {
 /// Table columns can be represented in two ways:
 ///
 /// 1. The table can have a column for each field in the `I` struct. In this case, the column names in the table
-/// should match the field names of the `I` struct. The field names can be a subset of column names and
-/// can be in any order.
+///    should match the field names of the `I` struct. The field names can be a subset of column names and
+///    can be in any order.
 ///
 /// For the supported field types, refer to:
 /// - [Serialization field types](https://docs.rs/serde_rusqlite/latest/serde_rusqlite)
 /// - [SQLite data types](https://www.sqlite.org/datatype3.html)
 ///
 /// 2. The fields in the `I` struct can be serialized into a single column `item` in the table. In this case, the table
-/// should have a single column named `item` of type `BLOB`. This is useful when the `I` struct contains complex fields
-/// that cannot be mapped to a SQLite type, such as nested structs, vectors, etc. The serialization is done using
-/// [MessagePack](https://msgpack.org/).
+///    should have a single column named `item` of type `BLOB`. This is useful when the `I` struct contains complex fields
+///    that cannot be mapped to a SQLite type, such as nested structs, vectors, etc. The serialization is done using
+///    [MessagePack](https://msgpack.org/).
 ///
 /// Note: The code automatically figures out which of the above two cases is applicable, and uses the appropriate
 /// method to read the data from the table.