Merge branch 'main' of github.com:cryscan/web-rwkv

cryscan · Nov 29, 2023 · 782343c · 782343c
2 parents 7af12d5 + bf43239
commit 782343c
Show file tree

Hide file tree

Showing 5 changed files with 33 additions and 32 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "web-rwkv"
-version = "0.4.0"
+version = "0.4.1"
 edition = "2021"
 authors = ["Zhenyuan Zhang <[email protected]>"]
 license = "MIT OR Apache-2.0"

diff --git a/src/model/mod.rs b/src/model/mod.rs
@@ -65,15 +65,15 @@ pub trait FromBuilder: Sized {
     fn from_builder(builder: Self::Builder<'_>) -> Result<Self, Self::Error>;
 }
 
-pub trait BackedState {
+pub trait BackedState: Send {
     fn max_batch(&self) -> usize;
     fn num_layer(&self) -> usize;
 
     /// Extract the embedding from a given layer of the state.
     fn embed(&self, batch: usize, layer: usize) -> Vec<f32>;
 }
 
-#[async_trait(?Send)]
+#[async_trait]
 pub trait ModelState {
     type BackedState: BackedState;
 
@@ -99,7 +99,7 @@ pub trait ModelState {
     ) -> Result<(), TensorError>;
 }
 
-#[async_trait(?Send)]
+#[async_trait]
 pub trait Model {
     type ModelState: ModelState;
 

diff --git a/src/model/v4.rs b/src/model/v4.rs
@@ -253,7 +253,7 @@ impl FromBuilder for ModelState {
     }
 }
 
-#[async_trait(?Send)]
+#[async_trait]
 impl super::ModelState for ModelState {
     type BackedState = BackedState;
 
@@ -272,7 +272,7 @@ impl super::ModelState for ModelState {
         if backed.max_batch() != self.max_batch() {
             return Err(ModelError::BatchSize(backed.max_batch(), self.max_batch()).into());
         }
-        let host = self.context.tensor_from_data(self.shape(), &backed.data)?;
+        let host = self.context.tensor_from_data(self.shape(), &*backed.data)?;
         self.0.load(&host).map_err(|err| err.into())
     }
 
@@ -283,7 +283,7 @@ impl super::ModelState for ModelState {
         }
         let shape = self.shape();
         let shape = Shape::new(shape[0], shape[1], 1, 1);
-        let host = self.context.tensor_from_data(shape, &backed.data)?;
+        let host = self.context.tensor_from_data(shape, &*backed.data)?;
         self.0.load_batch(&host, batch).map_err(|err| err.into())
     }
 
@@ -298,11 +298,8 @@ impl super::ModelState for ModelState {
         encoder.copy_tensor(self, &map).expect("back entire state");
         self.context.queue.submit(Some(encoder.finish()));
 
-        let host = map.back_async().await;
-        BackedState {
-            shape,
-            data: host.to_vec(),
-        }
+        let data = map.back_async().await.to_vec().into();
+        BackedState { shape, data }
     }
 
     async fn back_batch(&self, batch: usize) -> Result<Self::BackedState> {
@@ -325,11 +322,8 @@ impl super::ModelState for ModelState {
         encoder.copy_tensor_batch(self, &map, batch)?;
         self.context.queue.submit(Some(encoder.finish()));
 
-        let host = map.back_async().await;
-        Ok(BackedState {
-            shape,
-            data: host.to_vec(),
-        })
+        let data = map.back_async().await.to_vec().into();
+        Ok(BackedState { shape, data })
     }
 
     fn blit(&self, other: &Self) -> Result<(), TensorError> {
@@ -369,7 +363,7 @@ impl super::ModelState for ModelState {
 #[derive(Debug, Clone)]
 pub struct BackedState {
     pub shape: Shape,
-    pub data: Vec<f32>,
+    pub data: Arc<Vec<f32>>,
 }
 
 impl FromBuilder for BackedState {
@@ -398,7 +392,8 @@ impl FromBuilder for BackedState {
                     .concat()
             })
             .collect_vec()
-            .concat();
+            .concat()
+            .into();
         Ok(Self { shape, data })
     }
 }
@@ -910,7 +905,7 @@ impl<'a> FromBuilder for Model<'a> {
     }
 }
 
-#[async_trait(?Send)]
+#[async_trait]
 impl super::Model for Model<'_> {
     type ModelState = ModelState;
 

diff --git a/src/model/v5.rs b/src/model/v5.rs
@@ -272,8 +272,8 @@ impl FromBuilder for ModelState {
             })
             .collect();
         Ok(Self {
-            context: context.clone(),
-            info: info.clone(),
+            context,
+            info,
             max_batch,
             chunk_size,
             head_size,
@@ -282,7 +282,7 @@ impl FromBuilder for ModelState {
     }
 }
 
-#[async_trait(?Send)]
+#[async_trait]
 impl super::ModelState for ModelState {
     type BackedState = BackedState;
 
@@ -342,6 +342,7 @@ impl super::ModelState for ModelState {
             let host = map.back_async().await;
             data.push((shape, host.to_vec()))
         }
+        let data = data.into();
 
         BackedState {
             max_batch,
@@ -380,6 +381,7 @@ impl super::ModelState for ModelState {
             let host = map.back_async().await;
             data.push((shape, host.to_vec()));
         }
+        let data = data.into();
 
         Ok(BackedState {
             max_batch: 1,
@@ -433,7 +435,7 @@ pub struct BackedState {
     pub max_batch: usize,
     pub chunk_size: usize,
     pub head_size: usize,
-    pub data: Vec<(Shape, Vec<f32>)>,
+    pub data: Arc<Vec<(Shape, Vec<f32>)>>,
 }
 
 impl FromBuilder for BackedState {
@@ -457,7 +459,8 @@ impl FromBuilder for BackedState {
                     .concat()
             })
             .map(|x| (shape, x))
-            .collect();
+            .collect_vec()
+            .into();
         Ok(Self {
             max_batch,
             chunk_size,
@@ -1055,7 +1058,7 @@ impl<'a> FromBuilder for Model<'a> {
     }
 }
 
-#[async_trait(?Send)]
+#[async_trait]
 impl super::Model for Model<'_> {
     type ModelState = ModelState;
 

diff --git a/src/model/v6.rs b/src/model/v6.rs
@@ -307,8 +307,8 @@ impl FromBuilder for ModelState {
             })
             .collect();
         Ok(Self {
-            context: context.clone(),
-            info: info.clone(),
+            context,
+            info,
             max_batch,
             chunk_size,
             head_size,
@@ -317,7 +317,7 @@ impl FromBuilder for ModelState {
     }
 }
 
-#[async_trait(?Send)]
+#[async_trait]
 impl super::ModelState for ModelState {
     type BackedState = BackedState;
 
@@ -377,6 +377,7 @@ impl super::ModelState for ModelState {
             let host = map.back_async().await;
             data.push((shape, host.to_vec()))
         }
+        let data = data.into();
 
         BackedState {
             max_batch,
@@ -415,6 +416,7 @@ impl super::ModelState for ModelState {
             let host = map.back_async().await;
             data.push((shape, host.to_vec()));
         }
+        let data = data.into();
 
         Ok(BackedState {
             max_batch: 1,
@@ -468,7 +470,7 @@ pub struct BackedState {
     pub max_batch: usize,
     pub chunk_size: usize,
     pub head_size: usize,
-    pub data: Vec<(Shape, Vec<f32>)>,
+    pub data: Arc<Vec<(Shape, Vec<f32>)>>,
 }
 
 impl FromBuilder for BackedState {
@@ -492,7 +494,8 @@ impl FromBuilder for BackedState {
                     .concat()
             })
             .map(|x| (shape, x))
-            .collect();
+            .collect_vec()
+            .into();
         Ok(Self {
             max_batch,
             chunk_size,
@@ -1184,7 +1187,7 @@ impl<'a> FromBuilder for Model<'a> {
     }
 }
 
-#[async_trait(?Send)]
+#[async_trait]
 impl super::Model for Model<'_> {
     type ModelState = ModelState;