refactor(llama.cu): 剥离分批策略

Signed-off-by: YdrMaster <ydrml@hotmail.com>

Author: YdrMaster

llama.cu/src/batch/default.rs

@@ -0,0 +1,138 @@
+use super::{BatchStrategy, Req, Round, SessionId, SessionStub};
+use log::warn;
+use std::{cmp::min, collections::BTreeMap, iter::repeat_n, mem::take};
+
+pub(crate) struct DefaultStrategy<T> {
+    sess: BTreeMap<SessionId, SessionStub<T>>,
+    pre_output: BTreeMap<SessionId, usize>,
+    // 每次prefill的最大长度
+    chunked_prefill_max_len: Option<usize>,
+    max_toks: usize,
+}
+
+impl<T> DefaultStrategy<T> {
+    pub fn new(chunked_prefill_len: Option<usize>, max_toks: usize) -> Self {
+        Self {
+            sess: Default::default(),
+            pre_output: Default::default(),
+            chunked_prefill_max_len: chunked_prefill_len,
+            max_toks,
+        }
+    }
+}
+
+impl<T: 'static + Clone> BatchStrategy<T> for DefaultStrategy<T> {
+    fn is_empty(&self) -> bool {
+        self.sess.is_empty()
+    }
+
+    fn insert(&mut self, stub: SessionStub<T>) {
+        assert!(self.sess.insert(stub.session.id, stub).is_none())
+    }
+
+    fn remove(&mut self, id: &SessionId) -> Option<SessionStub<T>> {
+        self.sess.remove(id)
+    }
+
+    fn prepare(&mut self) -> Round<T> {
+        let mut ans = Round::default();
+        let mut out_idx = 0;
+
+        let pre_output = take(&mut self.pre_output);
+
+        let mut write_back_sessions = BTreeMap::new();
+
+        while let Some((id, mut stub)) = self.sess.pop_first() {
+            let max = stub.session.cache.capacity;
+            let pos = stub.session.cache.len;
+            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;
+            }
+
+            // 用于限制每次tokens总数
+            let remain_tok_num = self.max_toks - ans.tokens.len();
+            assert!(remain_tok_num > 0);
+
+            if let Some(prompt) = &stub.prompt {
+                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)
+            }
+
+            // 尝试填充缓存
+            stub.session.cache.len = end;
+            // 填充推理信息
+            ans.sample.extend(repeat_n(stub.session.sample_args, out));
+            ans.output.push((id, out));
+            ans.reqs.push(Req {
+                cache: stub.session.cache.cache.clone(),
+                pos,
+                seq,
+            });
+
+            //输出处理
+            //不会溢出 因为 out <= 1
+            stub.state.remain_steps -= out;
+            if stub.state.remain_steps == 0 {
+                // 生成结束
+                ans.finished.push(stub.session)
+            } else {
+                // 回填
+                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;
+            }
+        }
+        self.sess.append(&mut write_back_sessions);
+        ans
+    }
+
+    fn take_stubs(&mut self) -> Vec<SessionStub<T>> {
+        take(&mut self.sess).into_values().collect()
+    }
+}

llama.cu/src/batch/mod.rs

@@ -0,0 +1,75 @@
+mod default;
+
+use crate::SampleArgs;
+use tokeneer::utok;
+
+pub(crate) use default::DefaultStrategy;
+
+pub trait BatchStrategy<T: 'static> {
+    fn is_empty(&self) -> bool;
+    fn insert(&mut self, stub: SessionStub<T>);
+    fn remove(&mut self, id: &SessionId) -> Option<SessionStub<T>>;
+    fn prepare(&mut self) -> Round<T>;
+    fn take_stubs(&mut self) -> Vec<SessionStub<T>>;
+}
+
+// 目前在有prompt的情况下，state.seq 的长度代表prompt还有多少未prefill，也就是 `prompt[prompt.len() - state.seq..]` 代表未prefill的prompt
+pub(super) struct SessionStub<T> {
+    pub session: Session<T>,
+    pub state: State,
+    pub prompt: Option<Box<[utok]>>,
+}
+
+#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
+#[repr(transparent)]
+pub struct SessionId(pub usize);
+
+pub struct Round<T> {
+    pub overflow: Vec<Session<T>>,
+    pub tokens: Vec<utok>,
+    pub reqs: Vec<Req<T>>,
+    pub sample: Vec<SampleArgs>,
+    pub output: Vec<(SessionId, usize)>,
+    pub fast_map: Vec<(utok, utok)>,
+    pub finished: Vec<Session<T>>,
+}
+
+impl<T> Default for Round<T> {
+    fn default() -> Self {
+        Self {
+            overflow: Default::default(),
+            tokens: Default::default(),
+            reqs: Default::default(),
+            sample: Default::default(),
+            output: Default::default(),
+            fast_map: Default::default(),
+            finished: Default::default(),
+        }
+    }
+}
+
+pub struct Session<T> {
+    pub id: SessionId,
+    pub sample_args: SampleArgs,
+    pub cache: Cache<T>,
+}
+
+pub struct Cache<T> {
+    pub cache: T,
+    pub capacity: usize,
+    pub len: usize,
+}
+
+#[derive(Clone, Copy)]
+pub(super) struct State {
+    pub seq: usize,
+    pub out: usize,
+    pub remain_steps: usize,
+}
+
+#[derive(Clone)]
+pub(crate) struct Req<Cache> {
+    pub cache: Cache,
+    pub pos: usize,
+    pub seq: usize,
+}

llama.cu/src/exec/engine.rs

@@ -1,11 +1,10 @@
 use super::{
-    Command, Output, Request, Session,
-    engine_manager::{EngineManager, Round},
-    group::{ModelGroup, Req},
-    kv_cache::KVCache,
+    Command, Output, Request, engine_manager::EngineManager, group::ModelGroup,
     output_head::OutputHead,
 };
 use crate::{
+    CacheParts,
+    batch::{Req, Round, SessionStub, State},
     exec::{group::ModelGroupConfig, upos},
     handle::Handle,
     op::{FastEmbedding, random_sample::KVPair},
@@ -23,7 +22,7 @@ use std::{
     num::NonZeroUsize,
     ops::Deref,
     sync::{
-        Arc, Barrier, Mutex, OnceLock, RwLock,
+        Arc, Barrier, OnceLock, RwLock,
         atomic::AtomicUsize,
         mpsc::{Receiver, Sender},
     },
@@ -33,22 +32,10 @@ use tokeneer::utok;
 #[cfg(nccl)]
 use operators::nccl::{Communicator, CommunicatorGroup};
 
-// 目前在有prompt的情况下，state.seq 的长度代表prompt还有多少未prefill，也就是 `prompt[prompt.len() - state.seq..]` 代表未prefill的prompt
-pub(super) struct SessionStub {
-    pub session: Session,
-    pub state: State,
-    pub prompt: Option<Box<[utok]>>,
-}
-
-#[derive(Clone, Copy)]
-pub(super) struct State {
-    pub seq: usize,
-    pub out: usize,
-    pub remain_steps: usize,
-}
+type Stub = SessionStub<CacheParts>;
 
 impl Request {
-    pub(super) fn into_stub(self) -> SessionStub {
+    pub(super) fn into_stub(self) -> Stub {
         let Request {
             session,
             prompt,
@@ -206,7 +193,7 @@ type TaskBox = Arc<RwLock<Option<Task>>>;
 #[cfg_attr(not(nccl), allow(dead_code))]
 struct Task {
     key: NonZeroUsize,
-    reqs: Vec<Req<Arc<[Mutex<KVCache>]>>>,
+    reqs: Vec<Req<CacheParts>>,
 }
 
 impl<T: IntoIterator<Item = usize>> Worker<T> {