diff --git a/mistralrs-core/src/dummy_paged_attention/scheduler.rs b/mistralrs-core/src/dummy_paged_attention/scheduler.rs index 6aa92a180..9ccf9ebfa 100644 --- a/mistralrs-core/src/dummy_paged_attention/scheduler.rs +++ b/mistralrs-core/src/dummy_paged_attention/scheduler.rs @@ -345,7 +345,10 @@ impl Scheduler for PagedAttentionScheduler { } } fn waiting_len(&self) -> usize { - self.waiting.len() + self.waiting.len() + self.swapped_out.len() + } + fn running_len(&self) -> usize { + self.running.len() } fn block_tables(&self) -> Option<&BlockTables> { Some(&self.block_engine.block_tables) diff --git a/mistralrs-core/src/paged_attention/scheduler.rs b/mistralrs-core/src/paged_attention/scheduler.rs index 4e21cd70b..8e5bb8d2e 100644 --- a/mistralrs-core/src/paged_attention/scheduler.rs +++ b/mistralrs-core/src/paged_attention/scheduler.rs @@ -346,7 +346,10 @@ impl Scheduler for PagedAttentionScheduler { } } fn waiting_len(&self) -> usize { - self.waiting.len() + self.waiting.len() + self.swapped_out.len() + } + fn running_len(&self) -> usize { + self.running.len() } fn block_tables(&self) -> Option<&BlockTables> { Some(&self.block_engine.block_tables) diff --git a/mistralrs-core/src/scheduler/default_scheduler.rs b/mistralrs-core/src/scheduler/default_scheduler.rs index f36599b9a..22c543a8e 100644 --- a/mistralrs-core/src/scheduler/default_scheduler.rs +++ b/mistralrs-core/src/scheduler/default_scheduler.rs @@ -308,6 +308,9 @@ impl Scheduler for DefaultScheduler> { fn waiting_len(&self) -> usize { self.waiting.len() } + fn running_len(&self) -> usize { + self.running.len() + } fn add_seq(&mut self, seq: Sequence) { if seq.is_running() { // prefill case diff --git a/mistralrs-core/src/scheduler/mod.rs b/mistralrs-core/src/scheduler/mod.rs index 5f4b4fb5f..1f7302070 100644 --- a/mistralrs-core/src/scheduler/mod.rs +++ b/mistralrs-core/src/scheduler/mod.rs @@ -48,6 +48,7 @@ pub enum SchedulerOutput<'a> { pub trait Scheduler { fn schedule(&mut self) -> SchedulerOutput<'_>; fn waiting_len(&self) -> usize; + fn running_len(&self) -> usize; fn add_seq(&mut self, seq: Sequence); /// This may do nothing. It depends on the implementation fn free_finished_sequence_groups(&mut self);