Skip to content

Commit f09e3fd

Browse files
guoqingbaoguoqingbao
authored
Improve continuous batching (#419)
1 parent 1d23f69 commit f09e3fd

10 files changed

Lines changed: 204 additions & 143 deletions

File tree

‎Cargo.toml‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[package]
22
name = "candle-vllm"
3-
version = "0.7.5"
3+
version = "0.7.6"
44
edition = "2021"
55
default-run = "candle-vllm"
66

@@ -22,8 +22,8 @@ anyhow = "1.0.75"
2222
rand = "0.9.0"
2323
rayon="1.10.0"
2424
hyper = { version = "0.14", features = ["full"] }
25-
candle-core = { git = "https://github.com/guoqingbao/candle.git", version = "0.8.3", rev = "b63bf40" }
26-
candle-nn = { git = "https://github.com/guoqingbao/candle.git", version = "0.8.3", rev = "b63bf40" }
25+
candle-core = { git = "https://github.com/guoqingbao/candle.git", version = "0.8.3", rev = "68b6d74" }
26+
candle-nn = { git = "https://github.com/guoqingbao/candle.git", version = "0.8.3", rev = "68b6d74" }
2727
dyn-fmt = "0.4.0"
2828
safetensors = "0.4"
2929
serde = { version = "1.0.190", features = ["serde_derive"] }
@@ -46,7 +46,7 @@ dirs = "5.0.1"
4646
minijinja = { version = "2.10.2", features = ["builtins", "json"] }
4747
minijinja-contrib = { version = "2.10.2", features = ["pycompat"] }
4848
thiserror = "1.0.58"
49-
attention-rs = { git = "https://github.com/guoqingbao/attention.rs.git", version="0.5.2", rev = "2ea587f" }
49+
attention-rs = { git = "https://github.com/guoqingbao/attention.rs.git", version="0.5.3", rev = "dabe829" }
5050
metal = { version = "0.27.0", features = ["mps"], optional = true }
5151
lazy_static = {version = "1.4.0"}
5252
interprocess = "2.2.2"

‎src/api.rs‎

Lines changed: 11 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -444,14 +444,14 @@ impl Engine {
444444

445445
let (prompt, tokenizer, image_data, resolved_tools) = {
446446
let e = self.engine.read();
447-
let (pipeline, _) = e.get_pipeline(0).unwrap();
448447

449448
let tool_config = resolve_tools_for_request(&request.tools, &request.tool_choice, None)
450449
.map_err(candle_core::Error::wrap)?;
451450
let resolved_tools = tool_config.tools.clone();
452451

453-
// tokenizer is inside DefaultPipeline
454-
let mut conversation = pipeline.conversation.clone();
452+
let tokenizer = e.tokenizer().clone();
453+
let image_config = e.image_config();
454+
let mut conversation = e.conversation();
455455
let mut image_data = None;
456456

457457
// Logic to get prompt from messages
@@ -462,7 +462,7 @@ impl Engine {
462462
}
463463
Messages::Chat(messages) => {
464464
let (render_messages, images) =
465-
build_messages_and_images(messages, pipeline.image_config.as_ref())
465+
build_messages_and_images(messages, image_config.as_ref())
466466
.map_err(candle_core::Error::wrap)?;
467467
image_data = images;
468468
for message in render_messages {
@@ -499,12 +499,7 @@ impl Engine {
499499
let enable_thinking = request.thinking.unwrap_or(true);
500500
let prompt = conversation.get_prompt(enable_thinking, &tool_config.tools);
501501

502-
(
503-
prompt,
504-
pipeline.tokenizer.clone(),
505-
image_data,
506-
resolved_tools,
507-
)
502+
(prompt, tokenizer, image_data, resolved_tools)
508503
};
509504

510505
let request_id = format!("cmpl-{}", uuid::Uuid::new_v4());
@@ -692,10 +687,11 @@ impl Engine {
692687
}
693688

694689
let e = self.engine.read();
695-
let response_model = e
696-
.get_pipeline(0)
697-
.map(|(pipeline, _)| pipeline.name().to_string())
698-
.unwrap_or_else(|| request.model.clone().unwrap_or("default".to_string()));
690+
let response_model = if e.model_name().is_empty() {
691+
request.model.clone().unwrap_or("default".to_string())
692+
} else {
693+
e.model_name().to_string()
694+
};
699695
if let Some(record) = e.completion_records.get(&request_id) {
700696
let mut choices = record.0.clone();
701697
if crate::stream_as_reasoning_content() {
@@ -780,7 +776,6 @@ impl Engine {
780776
pub async fn embed_async(&self, request: EmbeddingRequest) -> Result<EmbeddingResponse> {
781777
let prompt_tokens = {
782778
let e = self.engine.read();
783-
let (pipeline, _) = e.get_pipeline(0).unwrap();
784779

785780
let prompt_str = match &request.input {
786781
crate::openai::requests::EmbeddingInput::String(s) => s.clone(),
@@ -799,8 +794,7 @@ impl Engine {
799794
}
800795
};
801796

802-
pipeline
803-
.tokenizer
797+
e.tokenizer()
804798
.encode(prompt_str, false)
805799
.map_err(candle_core::Error::msg)?
806800
.get_ids()

‎src/lib.rs‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -211,10 +211,7 @@ pub struct HybridMambaCachePlan {
211211
const DEFAULT_HYBRID_MAMBA_FRACTION: f32 = 0.15;
212212
const MAX_HYBRID_MAMBA_FRACTION: f32 = 0.3;
213213
const HYBRID_MAMBA_PREFIX_SLOT_MULTIPLIER: usize = 2;
214-
#[cfg(feature = "cuda")]
215214
const HYBRID_MAMBA_MIN_ACTIVE_SLOTS: usize = 8;
216-
#[cfg(not(feature = "cuda"))]
217-
const HYBRID_MAMBA_MIN_ACTIVE_SLOTS: usize = 4;
218215

219216
#[cfg_attr(not(any(feature = "cuda", feature = "metal")), allow(dead_code))]
220217
fn compute_kvcache_budget_bytes(free_bytes: usize, fraction: f32) -> Result<usize> {

‎src/openai/models/layers/others.rs‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -13,18 +13,18 @@ impl NormX {
1313
Self::Rms(_, dt) | Self::Layer(_, dt) => *dt,
1414
};
1515
let in_dtype = xs.dtype();
16-
let xs = if in_dtype != norm_dtype {
17-
xs.to_dtype(norm_dtype)?
18-
} else {
19-
xs.clone()
20-
};
21-
let out = match self {
22-
Self::Rms(norm, _) => norm.forward(&xs)?,
23-
Self::Layer(norm, _) => norm.forward(&xs)?,
24-
};
25-
if out.dtype() != in_dtype {
16+
if in_dtype != norm_dtype {
17+
let xs = xs.to_dtype(norm_dtype)?;
18+
let out = match self {
19+
Self::Rms(norm, _) => norm.forward(&xs)?,
20+
Self::Layer(norm, _) => norm.forward(&xs)?,
21+
};
2622
out.to_dtype(in_dtype)
2723
} else {
24+
let out = match self {
25+
Self::Rms(norm, _) => norm.forward(xs)?,
26+
Self::Layer(norm, _) => norm.forward(xs)?,
27+
};
2828
Ok(out)
2929
}
3030
}

‎src/openai/openai_server.rs‎

Lines changed: 15 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -34,10 +34,12 @@ const REQUEST_ADMISSION_DECODE_BUDGET_TOKENS: usize = 4096;
3434

3535
fn current_model_name(data: &OpenAIServerData) -> Result<String, APIError> {
3636
let model = data.model.read();
37-
let (pipeline, _) = model
38-
.get_pipeline(0)
39-
.ok_or(APIError::new("Missing pipeline".to_string()))?;
40-
Ok(pipeline.name().to_string())
37+
let model_name = model.model_name();
38+
if model_name.is_empty() {
39+
Err(APIError::new("Missing pipeline".to_string()))
40+
} else {
41+
Ok(model_name.to_string())
42+
}
4143
}
4244

4345
fn resolve_response_model_name(requested: Option<&str>, current: &str) -> String {
@@ -53,11 +55,10 @@ async fn get_gen_prompt(
5355
request: &ChatCompletionRequest,
5456
tool_config: &ResolvedToolConfig,
5557
) -> Result<(String, Option<ImageData>), APIError> {
56-
let mut model = data.model.write();
57-
let pipeline = model
58-
.get_mut_pipeline(0)
59-
.ok_or(APIError::new("Missing pipeline".to_string()))?;
60-
let mut conversation = pipeline.0.get_conversation().clone();
58+
let model = data.model.read();
59+
let mut conversation = model.conversation();
60+
let image_config = model.image_config();
61+
drop(model);
6162
let mut image_data = None;
6263

6364
match &request.messages {
@@ -66,7 +67,7 @@ async fn get_gen_prompt(
6667
}
6768
Messages::Chat(messages) => {
6869
let (render_messages, images) =
69-
build_messages_and_images(messages, pipeline.0.image_config.as_ref())
70+
build_messages_and_images(messages, image_config.as_ref())
7071
.map_err(APIError::from)?;
7172
image_data = images;
7273
for message in render_messages {
@@ -120,11 +121,7 @@ async fn check_length(
120121
) -> Result<Vec<u32>, APIError> {
121122
let token_ids = {
122123
let model = data.model.read();
123-
let pipeline = model
124-
.get_pipeline(0)
125-
.ok_or(APIError::new("Missing pipeline".to_string()))?;
126-
pipeline
127-
.0
124+
model
128125
.tokenizer()
129126
.encode_fast(prompt, true)
130127
.map_err(APIError::from)?
@@ -564,16 +561,9 @@ pub async fn create_embeddings(
564561
//TODO: Reuse check_length or similar logic. For now simplified.
565562
let token_ids = {
566563
let model = data.model.read();
567-
let pipeline = model
568-
.get_pipeline(0)
569-
.ok_or(APIError::new("Missing pipeline".to_string()));
570-
571-
match pipeline {
572-
Ok(pipeline) => match pipeline.0.tokenizer().encode_fast(prompt_str, true) {
573-
Ok(encoding) => encoding.get_ids().to_vec(),
574-
Err(e) => return ChatResponder::ValidationError(APIError::from(e)),
575-
},
576-
Err(e) => return ChatResponder::ModelError(e),
564+
match model.tokenizer().encode_fast(prompt_str, true) {
565+
Ok(encoding) => encoding.get_ids().to_vec(),
566+
Err(e) => return ChatResponder::ValidationError(APIError::from(e)),
577567
}
578568
};
579569

‎src/openai/pipelines/inputs.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,7 @@ impl LLMEngine {
140140
#[cfg(feature = "flashinfer")]
141141
prefill_tokens.push(num_tokens);
142142

143-
context_lens.push(seq_len as u32);
143+
context_lens.push((num_cached_tokens + num_tokens) as u32);
144144

145145
let seqlen_q = num_tokens;
146146
let use_cached_kv = num_cached_tokens > 0

0 commit comments

Comments
 (0)