refactor(transformer-cpu): decode 需要返回 batching logits 的顺序

Signed-off-by: YdrMaster <[email protected]>

service/src/cpu.rs

@@ -43,7 +43,7 @@ impl CpuTask {
                     .or_insert_with_key(|&id| SessionContext::new(&self.transformer, id));
 
                 let time = Instant::now();
-                let mut logits = self.transformer.decode(vec![ctx.request(&prompt)]);
+                let mut logits = self.transformer.decode(vec![ctx.request(&prompt)]).1;
                 info!("prefill transformer ... {:?}", time.elapsed());
 
                 let max_seq_len = self.transformer.max_seq_len() as upos;
@@ -54,7 +54,7 @@ impl CpuTask {
                     }
                     responsing.send(token).unwrap();
 
-                    logits = self.transformer.decode(vec![ctx.request(&[token])]);
+                    logits = self.transformer.decode(vec![ctx.request(&[token])]).1;
                 }
             }
             Command::Drop { id } => {
@@ -71,6 +71,7 @@ struct SessionContext {
 }
 
 impl SessionContext {
+    #[inline]
     fn new(transformer: &Transformer, id: usize) -> Self {
         Self {
             id,
@@ -79,6 +80,7 @@ impl SessionContext {
         }
     }
 
+    #[inline]
     fn request<'a>(&'a mut self, tokens: &'a [utok]) -> Request<usize> {
         let pos = self.pos;
         self.pos += tokens.len() as upos;

transformer-cpu/src/lib.rs

@@ -27,7 +27,7 @@ impl Transformer {
         self.0.max_position_embeddings()
     }
 
-    pub fn decode<Id>(&mut self, mut requests: Vec<Request<Id>>) -> Tensor<Storage> {
+    pub fn decode<Id>(&mut self, mut requests: Vec<Request<Id>>) -> (Vec<Id>, Tensor<Storage>) {
         requests.sort_unstable_by_key(|t| t.tokens.len());
 
         // println!("tokens:");
@@ -221,7 +221,7 @@ impl Transformer {
         mat_mul(&mut logits.access_mut(), 0., &x.access(), &lm_head, 1.);
         // println!("pos {pos} logits:\n{}", logits.access());
 
-        logits
+        (requests.into_iter().map(|r| r.id).collect(), logits)
     }
 }