Skip to content

Commit

Permalink
Nicer error when misconfigured pagedattn metadata (#753)
Browse files Browse the repository at this point in the history
  • Loading branch information
@EricLBuehler
EricLBuehler committed Sep 5, 2024
1 parent 9a9ec17 commit a625394
Show file tree
Hide file tree
Showing 3 changed files with 39 additions and 39 deletions.
46 changes: 21 additions & 25 deletions mistralrs-core/src/pipeline/gguf.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@ use crate::{
utils::tokens::get_token,
xlora_models::{XLoraQLlama, XLoraQPhi3},
};
use anyhow::{bail, Result};
use anyhow::{bail, Context, Result};
use candle_core::{DType, Device, Tensor};
use either::Either;
use hf_hub::{api::sync::ApiBuilder, Repo, RepoType};
Expand Down Expand Up @@ -645,29 +645,29 @@ impl Pipeline for GGUFPipeline {
flash_meta,
flash_meta_full,
} = *inputs.downcast().expect("Downcast failed.");
let paged_attn_meta = paged_attn_meta
.as_mut()
.with_context(|| "Forward step expected a PagedAttention input metadata. This was not provided, please ensure that the scheduler config is correctly configured for PagedAttention.")
.map_err(|e| candle_core::Error::Msg(e.to_string()))?;
let logits = match self.model {
Model::Llama(ref model) => model.forward(
&input_ids,
&seqlen_offsets,
seqlen_offsets_kernel,
context_lens,
self.get_metadata().cache_engine.as_ref().map(|engine| {
(
engine.get_kv_cache().clone(),
paged_attn_meta.as_mut().unwrap(),
)
}),
self.get_metadata()
.cache_engine
.as_ref()
.map(|engine| (engine.get_kv_cache().clone(), paged_attn_meta)),
)?,
Model::Phi2(ref model) => model.forward(
&input_ids,
&seqlen_offsets,
context_lens,
self.get_metadata().cache_engine.as_ref().map(|engine| {
(
engine.get_kv_cache().clone(),
paged_attn_meta.as_mut().unwrap(),
)
}),
self.get_metadata()
.cache_engine
.as_ref()
.map(|engine| (engine.get_kv_cache().clone(), paged_attn_meta)),
)?,
Model::XLoraLlama(ref model) => model.forward(
&input_ids,
Expand All @@ -685,12 +685,10 @@ impl Pipeline for GGUFPipeline {
Model::Phi3(ref model) => model.forward(
&input_ids,
&seqlen_offsets,
self.get_metadata().cache_engine.as_ref().map(|engine| {
(
engine.get_kv_cache().clone(),
paged_attn_meta.as_mut().unwrap(),
)
}),
self.get_metadata()
.cache_engine
.as_ref()
.map(|engine| (engine.get_kv_cache().clone(), paged_attn_meta)),
)?,
Model::XLoraPhi3(ref model) => model.forward(
&input_ids,
Expand All @@ -709,12 +707,10 @@ impl Pipeline for GGUFPipeline {
&input_ids,
&seqlen_offsets,
seqlen_offsets_kernel,
self.get_metadata().cache_engine.as_ref().map(|engine| {
(
engine.get_kv_cache().clone(),
paged_attn_meta.as_mut().unwrap(),
)
}),
self.get_metadata()
.cache_engine
.as_ref()
.map(|engine| (engine.get_kv_cache().clone(), paged_attn_meta)),
)?,
};
Ok(ForwardInputsResult::CausalGeneration { logits })
Expand Down
16 changes: 9 additions & 7 deletions mistralrs-core/src/pipeline/normal.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ use crate::{
normal_model_loader, xlora_model_loader, DeviceMapMetadata, PagedAttentionConfig, Pipeline,
Topology, TryIntoDType,
};
use anyhow::Result;
use anyhow::{Context, Result};
use candle_core::{Device, Tensor, Var};
use hf_hub::{api::sync::ApiBuilder, Repo, RepoType};
use mistralrs_quant::IsqType;
Expand Down Expand Up @@ -515,19 +515,21 @@ impl Pipeline for NormalPipeline {
flash_meta,
flash_meta_full,
} = *inputs.downcast().expect("Downcast failed.");
let paged_attn_meta = paged_attn_meta
.as_mut()
.with_context(|| "Forward step expected a PagedAttention input metadata. This was not provided, please ensure that the scheduler config is correctly configured for PagedAttention.")
.map_err(|e| candle_core::Error::Msg(e.to_string()))?;
let logits = match self.model.is_xlora() {
false => self.model.forward(
&input_ids,
&seqlen_offsets,
seqlen_offsets_kernel,
context_lens,
position_ids,
self.get_metadata().cache_engine.as_ref().map(|engine| {
(
engine.get_kv_cache().clone(),
paged_attn_meta.as_mut().unwrap(),
)
}),
self.get_metadata()
.cache_engine
.as_ref()
.map(|engine| (engine.get_kv_cache().clone(), paged_attn_meta)),
&flash_meta,
)?,
true => self.model.xlora_forward(
Expand Down
16 changes: 9 additions & 7 deletions mistralrs-core/src/pipeline/vision.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ use crate::{
api_dir_list, api_get_file, get_paths, vision_normal_model_loader, AnyMoeExpertType,
DeviceMapMetadata, Ordering, PagedAttentionConfig, Pipeline, Topology, TryIntoDType,
};
use anyhow::Result;
use anyhow::{Context, Result};
use candle_core::{Device, Tensor, Var};
use hf_hub::{api::sync::ApiBuilder, Repo, RepoType};
use mistralrs_quant::IsqType;
Expand Down Expand Up @@ -411,6 +411,10 @@ impl Pipeline for VisionPipeline {
mut paged_attn_meta,
flash_meta,
} = *inputs.downcast::<ModelInputs>().expect("Downcast failed.");
let paged_attn_meta = paged_attn_meta
.as_mut()
.with_context(|| "Forward step expected a PagedAttention input metadata. This was not provided, please ensure that the scheduler config is correctly configured for PagedAttention.")
.map_err(|e| candle_core::Error::Msg(e.to_string()))?;
let logits = self.model.forward(
&input_ids,
pixel_values,
Expand All @@ -419,12 +423,10 @@ impl Pipeline for VisionPipeline {
context_lens,
position_ids,
model_specific_args,
self.get_metadata().cache_engine.as_ref().map(|engine| {
(
engine.get_kv_cache().clone(),
paged_attn_meta.as_mut().unwrap(),
)
}),
self.get_metadata()
.cache_engine
.as_ref()
.map(|engine| (engine.get_kv_cache().clone(), paged_attn_meta)),
&flash_meta,
)?;
Ok(ForwardInputsResult::CausalGeneration { logits })
Expand Down

0 comments on commit a625394

Please sign in to comment.