Skip to content

Commit 5cdaee0

Browse files
alvarobarttmichaelfeil
alvarobartt
and
michaelfeil
authored
Add use_flash_attn for better FA + FA2 feature gating (#825)
Co-authored-by: Michael Feil <michaelfeil@users.noreply.github.com>
1 parent ff2f0cc commit 5cdaee0

1 file changed

Lines changed: 47 additions & 67 deletions

File tree

‎backends/candle/src/lib.rs‎

Lines changed: 47 additions & 67 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,38 @@ use crate::models::{
3535
FlashQwen2Model, FlashQwen3Model,
3636
};
3737

38+
#[derive(Debug, Clone, Copy, PartialEq)]
39+
pub enum FlashAttn {
40+
V1,
41+
V2,
42+
}
43+
44+
#[allow(unused)]
45+
fn use_flash_attn(supported: &[FlashAttn]) -> bool {
46+
#[cfg(not(feature = "cuda"))]
47+
{
48+
tracing::warn!("Flash Attention is not supported on CPU yet");
49+
false
50+
}
51+
#[cfg(feature = "cuda")]
52+
{
53+
if std::env::var("USE_FLASH_ATTENTION")
54+
.unwrap_or_else(|_| "True".to_string())
55+
.to_lowercase()
56+
!= "true"
57+
{
58+
false
59+
} else {
60+
supported.iter().any(|v| match v {
61+
FlashAttn::V1 => cfg!(feature = "flash-attn-v1"),
62+
FlashAttn::V2 => {
63+
cfg!(feature = "flash-attn") && get_runtime_compute_cap().is_ok_and(|x| x >= 80)
64+
}
65+
})
66+
}
67+
}
68+
}
69+
3870
/// This enum is needed to be able to differentiate between jina models that also use
3971
/// the `bert` model type and valid Bert models.
4072
#[derive(Debug, Clone, PartialEq)]
@@ -332,12 +364,7 @@ impl CandleBackend {
332364
}
333365
#[cfg(feature = "cuda")]
334366
(Config::Bert(config), Device::Cuda(_)) => {
335-
if cfg!(any(feature = "flash-attn", feature = "flash-attn-v1"))
336-
&& dtype == DType::F16
337-
// Allow disabling because of flash attention v1 precision problems
338-
// See: https://github.com/huggingface/text-embeddings-inference/issues/37
339-
&& &std::env::var("USE_FLASH_ATTENTION").unwrap_or("True".to_string()).to_lowercase() == "true"
340-
{
367+
if dtype == DType::F16 && use_flash_attn(&[FlashAttn::V1, FlashAttn::V2]) {
341368
match config {
342369
BertConfigWrapper::JinaBert(config) => {
343370
tracing::info!("Starting FlashJinaBert model on {:?}", device);
@@ -380,12 +407,7 @@ impl CandleBackend {
380407
Config::Camembert(config) | Config::Roberta(config) | Config::XlmRoberta(config),
381408
Device::Cuda(_),
382409
) => {
383-
if cfg!(any(feature = "flash-attn", feature = "flash-attn-v1"))
384-
&& dtype == DType::F16
385-
// Allow disabling because of flash attention v1 precision problems
386-
// See: https://github.com/huggingface/text-embeddings-inference/issues/37
387-
&& &std::env::var("USE_FLASH_ATTENTION").unwrap_or("True".to_string()).to_lowercase() == "true"
388-
{
410+
if dtype == DType::F16 && use_flash_attn(&[FlashAttn::V1, FlashAttn::V2]) {
389411
tracing::info!("Starting FlashBert model on {:?}", device);
390412
Ok(Box::new(
391413
FlashBertModel::load_roberta(vb, &config, model_type).s()?,
@@ -404,13 +426,7 @@ impl CandleBackend {
404426
}
405427
#[cfg(feature = "cuda")]
406428
(Config::DistilBert(config), Device::Cuda(_)) => {
407-
if cfg!(feature = "flash-attn")
408-
&& dtype == DType::F16
409-
&& &std::env::var("USE_FLASH_ATTENTION")
410-
.unwrap_or("True".to_string())
411-
.to_lowercase()
412-
== "true"
413-
{
429+
if dtype == DType::F16 && use_flash_attn(&[FlashAttn::V2]) {
414430
tracing::info!("Starting FlashDistilBert model on {:?}", device);
415431
Ok(Box::new(
416432
FlashDistilBertModel::load(vb, &config, model_type).s()?,
@@ -435,30 +451,17 @@ impl CandleBackend {
435451
}
436452
#[cfg(feature = "cuda")]
437453
(Config::Gte(config), Device::Cuda(_)) => {
438-
if dtype != DType::F16
439-
|| !cfg!(any(feature = "flash-attn", feature = "flash-attn-v1"))
440-
|| &std::env::var("USE_FLASH_ATTENTION")
441-
.unwrap_or("True".to_string())
442-
.to_lowercase()
443-
!= "true"
444-
{
445-
tracing::info!("Starting GTE model on {:?}", device);
446-
Ok(Box::new(GTEModel::load(vb, &config, model_type).s()?))
447-
} else {
454+
if dtype == DType::F16 && use_flash_attn(&[FlashAttn::V1, FlashAttn::V2]) {
448455
tracing::info!("Starting FlashGTE model on {:?}", device);
449456
Ok(Box::new(FlashGTEModel::load(vb, &config, model_type).s()?))
457+
} else {
458+
tracing::info!("Starting GTE model on {:?}", device);
459+
Ok(Box::new(GTEModel::load(vb, &config, model_type).s()?))
450460
}
451461
}
452462
#[cfg(feature = "cuda")]
453463
(Config::Mistral(config), Device::Cuda(_)) => {
454-
if dtype != DType::F16
455-
|| !cfg!(feature = "flash-attn")
456-
|| get_runtime_compute_cap().unwrap() < 80
457-
|| &std::env::var("USE_FLASH_ATTENTION")
458-
.unwrap_or("True".to_string())
459-
.to_lowercase()
460-
!= "true"
461-
{
464+
if !(dtype == DType::F16 && use_flash_attn(&[FlashAttn::V2])) {
462465
return Err(BackendError::Start("Mistral is only supported on Cuda devices in fp16 with flash attention v2 enabled".to_string()));
463466
}
464467
tracing::info!("Starting FlashMistral model on {:?}", device);
@@ -468,12 +471,7 @@ impl CandleBackend {
468471
}
469472
#[cfg(feature = "cuda")]
470473
(Config::ModernBert(config), Device::Cuda(_)) => {
471-
if cfg!(feature = "flash-attn")
472-
&& dtype == DType::F16
473-
// Allow disabling because of flash attention v1 precision problems
474-
// See: https://github.com/huggingface/text-embeddings-inference/issues/37
475-
&& &std::env::var("USE_FLASH_ATTENTION").unwrap_or("True".to_string()).to_lowercase() == "true"
476-
{
474+
if dtype == DType::F16 && use_flash_attn(&[FlashAttn::V2]) {
477475
tracing::info!("Starting FlashModernBert model on {:?}", device);
478476
Ok(Box::new(
479477
FlashModernBertModel::load(vb, &config, model_type).s()?,
@@ -489,13 +487,7 @@ impl CandleBackend {
489487
}
490488
#[cfg(feature = "cuda")]
491489
(Config::NomicBert(config), Device::Cuda(_)) => {
492-
if cfg!(feature = "flash-attn")
493-
&& dtype == DType::F16
494-
&& &std::env::var("USE_FLASH_ATTENTION")
495-
.unwrap_or("True".to_string())
496-
.to_lowercase()
497-
== "true"
498-
{
490+
if dtype == DType::F16 && use_flash_attn(&[FlashAttn::V2]) {
499491
tracing::info!("Starting FlashNomicBert model on {:?}", device);
500492
Ok(Box::new(
501493
FlashNomicBertModel::load(vb, &config, model_type).s()?,
@@ -507,13 +499,7 @@ impl CandleBackend {
507499
}
508500
#[cfg(feature = "cuda")]
509501
(Config::Qwen2(config), Device::Cuda(_)) => {
510-
if dtype != DType::F16
511-
|| !cfg!(any(feature = "flash-attn", feature = "flash-attn-v1"))
512-
|| &std::env::var("USE_FLASH_ATTENTION")
513-
.unwrap_or("True".to_string())
514-
.to_lowercase()
515-
!= "true"
516-
{
502+
if !(dtype == DType::F16 && use_flash_attn(&[FlashAttn::V1, FlashAttn::V2])) {
517503
return Err(BackendError::Start("Qwen2 is only supported on Cuda devices in fp16 with flash attention v2 enabled".to_string()));
518504
}
519505
tracing::info!("Starting FlashQwen2 model on {:?}", device);
@@ -523,20 +509,14 @@ impl CandleBackend {
523509
}
524510
#[cfg(feature = "cuda")]
525511
(Config::Qwen3(config), Device::Cuda(_)) => {
526-
if dtype != DType::F16
527-
|| !cfg!(any(feature = "flash-attn", feature = "flash-attn-v1"))
528-
|| &std::env::var("USE_FLASH_ATTENTION")
529-
.unwrap_or("True".to_string())
530-
.to_lowercase()
531-
!= "true"
532-
{
533-
tracing::info!("Starting Qwen3 model on {:?}", device);
534-
Ok(Box::new(Qwen3Model::load(vb, &config, model_type).s()?))
535-
} else {
512+
if dtype == DType::F16 && use_flash_attn(&[FlashAttn::V1, FlashAttn::V2]) {
536513
tracing::info!("Starting FlashQwen3 model on {:?}", device);
537514
Ok(Box::new(
538515
FlashQwen3Model::load(vb, &config, model_type).s()?,
539516
))
517+
} else {
518+
tracing::info!("Starting Qwen3 model on {:?}", device);
519+
Ok(Box::new(Qwen3Model::load(vb, &config, model_type).s()?))
540520
}
541521
}
542522
#[cfg(feature = "cuda")]

0 commit comments

Comments
 (0)