Merge pull request #37 from pwhMass/fix_chunked_prefill

fix: Fix chunked prefill

Author: YdrMaster

llama.cu/src/exec/engine.rs

@@ -66,9 +66,14 @@ impl Request {
     }
 }
 
+// 用于测试chunked prefill和动态建图
+// const NTOKS: [usize; 2] = [1, 4];
+// const CHUNKED_PREFILL_LEN: Option<usize> = Some(2);
+// const MAX_TOKS: usize = 8;
+
+// TODO这些参数可能需要优化，目前是根据经验设置的
 const NTOKS: [usize; 7] = [1, 8, 32, 64, 128, 256, 1024];
-const CHUNKED_PREFILL_LEN: Option<usize> = Some(32);
-//TODO 该常量应该放在哪比较合适
+const CHUNKED_PREFILL_LEN: Option<usize> = Some(256);
 const MAX_TOKS: usize = 1024;
 
 pub(crate) fn engine(
@@ -203,7 +208,7 @@ impl<T: IntoIterator<Item = usize>> Worker<T> {
         let gpu = Gpu::new(dev.retain_primary(), Default::default());
         let attn = Attn::new(&gpu);
         gpu.apply(|ctx| {
-            let mut manager = EngineManager::new(chunked_prefill_len);
+            let mut manager = EngineManager::new(chunked_prefill_len, max_toks);
             let mut handle = handle(ctx);
             let mut models =
                 ModelGroup::new(llama, dist, config, attn, &mut handle, barrier.as_deref());
@@ -283,18 +288,10 @@ impl<T: IntoIterator<Item = usize>> Worker<T> {
 
                     // 如果没有输出，则跳过
                     if !out_idx.is_empty() {
-                        let (output, sample): (Vec<_>, Vec<_>) = output
-                            .iter()
-                            .zip(sample.iter())
-                            .filter_map(|((id, len), sample_arg)| {
-                                if *len > 0 {
-                                    Some(((*id, *len), sample_arg))
-                                } else {
-                                    None
-                                }
-                            })
-                            .unzip();
-
+                        let output = output
+                            .into_iter()
+                            .filter_map(|(id, len)| if len > 0 { Some((id, len)) } else { None })
+                            .collect::<Vec<_>>();
                         let kv_pairs = output_head.launch(
                             x,
                             &out_idx_buf[..out_idx.len()],

llama.cu/src/exec/engine_manager.rs

@@ -2,6 +2,7 @@
 use crate::{exec::KVCache, op::random_sample::SampleArgs};
 use log::warn;
 use std::{
+    cmp::min,
     collections::BTreeMap,
     iter::repeat_n,
     mem::take,
@@ -12,12 +13,12 @@ use std::{
 };
 use tokeneer::utok;
 
-#[derive(Default)]
 pub(super) struct EngineManager {
     sess: BTreeMap<SessionId, SessionStub>,
     pre_output: BTreeMap<SessionId, usize>,
     // 每次prefill的最大长度
-    chunked_prefill_len: Option<usize>,
+    chunked_prefill_max_len: Option<usize>,
+    max_toks: usize,
 }
 
 #[derive(Default)]
@@ -41,10 +42,12 @@ pub enum CommandError {
 type E = CommandError;
 
 impl EngineManager {
-    pub fn new(chunked_prefill_len: Option<usize>) -> Self {
+    pub fn new(chunked_prefill_len: Option<usize>, max_toks: usize) -> Self {
         Self {
-            chunked_prefill_len,
-            ..Default::default()
+            sess: Default::default(),
+            pre_output: Default::default(),
+            chunked_prefill_max_len: chunked_prefill_len,
+            max_toks,
         }
     }
     /// 接收命令
@@ -82,53 +85,62 @@ impl EngineManager {
         let mut out_idx = 0;
 
         let pre_output = take(&mut self.pre_output);
-        for (id, mut stub) in take(&mut self.sess) {
+
+        let mut write_back_sessions = BTreeMap::<SessionId, SessionStub>::new();
+
+        while let Some((id, mut stub)) = self.sess.pop_first() {
             let max = stub.session.cache.len;
             let pos = stub.session.cache.pos;
-            let seq = stub.state.seq;
-            let out = stub.state.out;
-            let end = pos + seq;
-            assert_eq!(out, 1, "TODO: ???");
+            let mut seq = stub.state.seq;
+            let mut out = stub.state.out;
+            let mut end = pos + seq;
+            assert_eq!(out, 1, "TODO: 投机采样");
             //验证缓存是否溢出
             if end > max {
                 warn!("cache overflow {end} > {max}");
                 // 缓存溢出，不再推理
                 ans.overflow.push(stub.session);
                 continue;
             }
-            //chunked prefill
-            // 采用 state.seq 用于计算剩余需要prefill的长度
+
+            // 用于限制每次tokens总数
+            let remain_tok_num = self.max_toks - ans.tokens.len();
+            assert!(remain_tok_num > 0);
+
             if let Some(prompt) = &stub.prompt {
-                if let Some(chunked_prefill_len) = self.chunked_prefill_len {
-                    if stub.state.seq > chunked_prefill_len {
-                        // 根据chunked_prefill_len重新计算seq和end
-                        let seq = chunked_prefill_len;
-                        let end = pos + seq;
-                        ans.sample.push(stub.session.sample_args);
-                        ans.output.push((id, 0));
-                        ans.reqs.push(Req {
-                            kv_cache: stub.session.cache.parts.clone(),
-                            pos,
-                            seq,
-                        });
-                        ans.tokens.extend(
-                            prompt
-                                .iter()
-                                .skip(prompt.len() - stub.state.seq)
-                                .take(chunked_prefill_len),
-                        );
-
-                        //更新stub信息
-                        stub.session.cache.pos = end;
-                        stub.state.seq -= chunked_prefill_len;
-
-                        //回填
-                        assert!(self.sess.insert(id, stub).is_none());
-
-                        //提前结束
-                        continue;
+                seq = self
+                    .chunked_prefill_max_len
+                    .map_or(min(remain_tok_num, seq), |chunked_prefill_max_len| {
+                        remain_tok_num.min(seq).min(chunked_prefill_max_len)
+                    });
+
+                if seq < stub.state.seq {
+                    // chunked prefill
+                    out = 0;
+                    end = pos + seq;
+
+                    ans.tokens
+                        .extend(prompt.iter().skip(prompt.len() - stub.state.seq).take(seq));
+
+                    //更新stub信息
+                    stub.state.seq -= seq;
+                } else {
+                    // 正常prefill
+                    if seq != prompt.len() {
+                        log::debug!("{:?} chunked prefil finished", id);
                     }
+                    ans.tokens.extend(prompt[prompt.len() - seq..].to_owned());
+
+                    stub.state.seq = 1;
+                    stub.prompt = None;
                 }
+            } else {
+                // decode
+                assert_eq!(seq, 1);
+                // fast embd
+                ans.fast_map
+                    .push((pre_output[&id] as _, ans.tokens.len() as _));
+                ans.tokens.push(0)
             }
 
             // 尝试填充缓存
@@ -141,35 +153,28 @@ impl EngineManager {
                 pos,
                 seq,
             });
-            if let Some(prompt) = stub.prompt.take() {
-                // prefill
-                if seq != prompt.len() {
-                    log::debug!("{:?} chunked prefil finished", id);
-                }
-                ans.tokens.extend(prompt[prompt.len() - seq..].to_owned());
-
-                stub.state.seq = 1
-            } else {
-                // decode
-                assert_eq!(seq, 1);
-                // fast embd
-                ans.fast_map
-                    .push((pre_output[&id] as _, ans.tokens.len() as _));
-                ans.tokens.push(0)
-            }
 
             //输出处理
-            stub.state.remain_steps -= 1;
+            //不会溢出 因为 out <= 1
+            stub.state.remain_steps -= out;
             if stub.state.remain_steps == 0 {
                 // 生成结束
                 ans.finished.push(stub.session)
             } else {
                 // 回填
-                assert!(self.sess.insert(id, stub).is_none());
-                assert!(self.pre_output.insert(id, out_idx).is_none());
+                assert!(write_back_sessions.insert(id, stub).is_none());
+                if out != 0 {
+                    assert!(self.pre_output.insert(id, out_idx).is_none());
+                }
+            }
+            out_idx += out;
+
+            // 如果剩余tokens总数等于0，则退出循环
+            if self.max_toks == ans.tokens.len() {
+                break;
             }
-            out_idx += out
         }
+        self.sess.append(&mut write_back_sessions);
         ans
     }
 

xtask/src/service/mod.rs

@@ -131,13 +131,15 @@ async fn start_infer_service(
                 } else {
                     &[]
                 };
-                if think.len() < tokens.len() {
-                    session_info.think = false;
-                    tokens = &tokens[think.len() + 1..]
-                } else {
-                    tokens = &[]
-                }
 
+                if session_info.think {
+                    if think.len() < tokens.len() {
+                        session_info.think = false;
+                        tokens = &tokens[think.len() + 1..]
+                    } else {
+                        tokens = &[]
+                    }
+                }
                 let think = service_manager_for_recv
                     .terminal
                     .decode(think, &mut session_info.buf);