add Qwen3VL model

This commit is contained in:
jhqxxx
2025-10-26 21:39:23 +08:00
parent cdae7fc45f
commit 7b190e7d54
22 changed files with 2591 additions and 53 deletions
+51
View File
@@ -269,3 +269,54 @@ impl AttentionNobias {
self.kv_cache = None
}
}
pub fn eager_attention_forward(
query_states: &Tensor,
key_states: &Tensor,
value_states: &Tensor,
num_key_value_groups: Option<usize>,
attention_mask: Option<&Tensor>,
scaling: f64,
) -> Result<Tensor> {
let key_states = match num_key_value_groups {
Some(g) => repeat_kv(key_states.clone(), g)?.contiguous()?,
None => key_states.clone()
};
let value_states = match num_key_value_groups {
Some(g) => repeat_kv(value_states.clone(), g)?.contiguous()?,
None => value_states.clone()
};
let attn_output = {
#[cfg(not(feature = "flash-attn"))]
{
let attn_weights = query_states.matmul(&key_states.transpose(D::Minus2, D::Minus1)?)?;
let attn_weights = (attn_weights * scaling)?;
let attn_weights = match attention_mask {
None => attn_weights,
Some(mask) => attn_weights.broadcast_add(&mask.to_dtype(attn_weights.dtype())?)?,
};
let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?;
attn_weights.matmul(&value_states)?
}
#[cfg(feature = "flash-attn")]
{
// use flash-attn,
// flash-attn shape: (bs, seq_len, num_head, head_dim)
let query_states = query_states.transpose(1, 2)?;
let key_states = key_states.transpose(1, 2)?;
let value_states = value_states.transpose(1, 2)?;
let attn_output = candle_flash_attn::flash_attn(
&query_states,
&key_states,
&value_states,
scaling as f32,
attention_mask.is_some(),
)?
.transpose(1, 2)?;
attn_output
}
};
let attn_output = attn_output.transpose(1, 2)?.contiguous()?;
Ok(attn_output)
}