diff --git a/candle-examples/examples/quantized-qwen3/README.md b/candle-examples/examples/quantized-qwen3/README.md new file mode 100644 index 00000000..2260536c --- /dev/null +++ b/candle-examples/examples/quantized-qwen3/README.md @@ -0,0 +1,11 @@ +# candle-quantized-qwen3 + +[Qwen3]((https://qwenlm.github.io/blog/qwen3/)) is an upgraded version of Qwen2.5, released by Alibaba Cloud. + +## Running the example + +```bash +cargo run --example quantized-qwen3 --release -- --prompt "Write a function to count prime numbers up to N." +``` + +0.6b is used by default, 1.7b, 4b, 8b, 14b, and 32b models are available via `--model` argument. diff --git a/candle-examples/examples/quantized-qwen3/main.rs b/candle-examples/examples/quantized-qwen3/main.rs new file mode 100644 index 00000000..b57466be --- /dev/null +++ b/candle-examples/examples/quantized-qwen3/main.rs @@ -0,0 +1,314 @@ +#[cfg(feature = "mkl")] +extern crate intel_mkl_src; + +#[cfg(feature = "accelerate")] +extern crate accelerate_src; + +use clap::{Parser, ValueEnum}; +use std::io::Write; +use tokenizers::Tokenizer; + +use candle::quantized::gguf_file; +use candle::Tensor; +use candle_transformers::generation::{LogitsProcessor, Sampling}; + +use candle_examples::token_output_stream::TokenOutputStream; +use candle_transformers::models::quantized_qwen3::ModelWeights as Qwen3; + +const DEFAULT_PROMPT: &str = "Write a Rust function to calculate the factorial of a given number."; + +#[derive(Clone, Debug, Copy, PartialEq, Eq, ValueEnum)] +enum Which { + #[value(name = "0.6b")] + W3_0_6b, + #[value(name = "1.7b")] + W3_1_7b, + #[value(name = "4b")] + W3_4b, + #[value(name = "8b")] + W3_8b, + #[value(name = "14b")] + W3_14b, + #[value(name = "32b")] + W3_32b, +} + +#[derive(Parser, Debug)] +#[command(author, version, about, long_about = None)] +struct Args { + /// GGUF file to load, typically a .gguf file generated by the quantize command from llama.cpp + #[arg(long)] + model: Option, + + /// The initial prompt, use 'interactive' for entering multiple prompts in an interactive way + /// and 'chat' for an interactive model where history of previous prompts and generated tokens + /// is preserved. + #[arg(long)] + prompt: Option, + + /// The length of the sample to generate (in tokens). + #[arg(short = 'n', long, default_value_t = 1000)] + sample_len: usize, + + /// The tokenizer config in json format. + #[arg(long)] + tokenizer: Option, + + /// The temperature used to generate samples, use 0 for greedy sampling. + #[arg(long, default_value_t = 0.8)] + temperature: f64, + + /// Nucleus sampling probability cutoff. + #[arg(long)] + top_p: Option, + + /// Only sample among the top K samples. + #[arg(long)] + top_k: Option, + + /// The seed to use when generating random samples. + #[arg(long, default_value_t = 299792458)] + seed: u64, + + /// Enable tracing (generates a trace-timestamp.json file). + #[arg(long)] + tracing: bool, + + /// Process prompt elements separately. + #[arg(long)] + split_prompt: bool, + + /// Run on CPU rather than GPU even if a GPU is available. + #[arg(long)] + cpu: bool, + + /// Penalty to be applied for repeating tokens, 1. means no penalty. + #[arg(long, default_value_t = 1.1)] + repeat_penalty: f32, + + /// The context size to consider for the repeat penalty. + #[arg(long, default_value_t = 64)] + repeat_last_n: usize, + + /// The model size to use. + #[arg(long, default_value = "0.6b")] + which: Which, +} + +impl Args { + fn tokenizer(&self) -> anyhow::Result { + let tokenizer_path = match &self.tokenizer { + Some(config) => std::path::PathBuf::from(config), + None => { + let api = hf_hub::api::sync::Api::new()?; + let repo = match self.which { + Which::W3_0_6b => "Qwen/Qwen3-0.6B", + Which::W3_1_7b => "Qwen/Qwen3-1.7B", + Which::W3_4b => "Qwen/Qwen3-4B", + Which::W3_8b => "Qwen/Qwen3-8B", + Which::W3_14b => "Qwen/Qwen3-14B", + Which::W3_32b => "Qwen/Qwen3-32B", + }; + let api = api.model(repo.to_string()); + api.get("tokenizer.json")? + } + }; + Tokenizer::from_file(tokenizer_path).map_err(anyhow::Error::msg) + } + + fn model(&self) -> anyhow::Result { + let model_path = match &self.model { + Some(config) => std::path::PathBuf::from(config), + None => { + let (repo, filename, revision) = match self.which { + Which::W3_0_6b => ("unsloth/Qwen3-0.6B-GGUF", "Qwen3-0.6B-Q4_K_M.gguf", "main"), + Which::W3_1_7b => ("unsloth/Qwen3-1.7B-GGUF", "Qwen3-1.7B-Q4_K_M.gguf", "main"), + Which::W3_4b => ("unsloth/Qwen3-4B-GGUF", "Qwen3-4B-Q4_K_M.gguf", "main"), + Which::W3_8b => ("unsloth/Qwen3-8B-GGUF", "Qwen3-8B-Q4_K_M.gguf", "main"), + Which::W3_14b => ("unsloth/Qwen3-14B-GGUF", "Qwen3-14B-Q4_K_M.gguf", "main"), + Which::W3_32b => ("unsloth/Qwen3-32B-GGUF", "Qwen3-32B-Q4_K_M.gguf", "main"), + }; + let api = hf_hub::api::sync::Api::new()?; + api.repo(hf_hub::Repo::with_revision( + repo.to_string(), + hf_hub::RepoType::Model, + revision.to_string(), + )) + .get(filename)? + } + }; + Ok(model_path) + } +} + +fn format_size(size_in_bytes: usize) -> String { + if size_in_bytes < 1_000 { + format!("{}B", size_in_bytes) + } else if size_in_bytes < 1_000_000 { + format!("{:.2}KB", size_in_bytes as f64 / 1e3) + } else if size_in_bytes < 1_000_000_000 { + format!("{:.2}MB", size_in_bytes as f64 / 1e6) + } else { + format!("{:.2}GB", size_in_bytes as f64 / 1e9) + } +} + +fn main() -> anyhow::Result<()> { + use tracing_chrome::ChromeLayerBuilder; + use tracing_subscriber::prelude::*; + + let args = Args::parse(); + let _guard = if args.tracing { + let (chrome_layer, guard) = ChromeLayerBuilder::new().build(); + tracing_subscriber::registry().with(chrome_layer).init(); + Some(guard) + } else { + None + }; + + println!( + "avx: {}, neon: {}, simd128: {}, f16c: {}", + candle::utils::with_avx(), + candle::utils::with_neon(), + candle::utils::with_simd128(), + candle::utils::with_f16c() + ); + println!( + "temp: {:.2} repeat-penalty: {:.2} repeat-last-n: {}", + args.temperature, args.repeat_penalty, args.repeat_last_n + ); + + let model_path = args.model()?; + let mut file = std::fs::File::open(&model_path)?; + let start = std::time::Instant::now(); + let device = candle_examples::device(args.cpu)?; + + let mut model = { + let model = gguf_file::Content::read(&mut file).map_err(|e| e.with_path(model_path))?; + let mut total_size_in_bytes = 0; + for (_, tensor) in model.tensor_infos.iter() { + let elem_count = tensor.shape.elem_count(); + total_size_in_bytes += + elem_count * tensor.ggml_dtype.type_size() / tensor.ggml_dtype.block_size(); + } + println!( + "loaded {:?} tensors ({}) in {:.2}s", + model.tensor_infos.len(), + &format_size(total_size_in_bytes), + start.elapsed().as_secs_f32(), + ); + Qwen3::from_gguf(model, &mut file, &device)? + }; + println!("model built"); + + let tokenizer = args.tokenizer()?; + let mut tos = TokenOutputStream::new(tokenizer); + let prompt_str = args + .prompt + .clone() + .unwrap_or_else(|| DEFAULT_PROMPT.to_string()); + + let prompt_str = format!("<|im_start|>user\n{prompt_str}<|im_end|>\n<|im_start|>assistant\n"); + print!("formatted prompt: {}", &prompt_str); + + let tokens = tos + .tokenizer() + .encode(prompt_str, true) + .map_err(anyhow::Error::msg)?; + + let tokens = tokens.get_ids(); + + let to_sample = args.sample_len.saturating_sub(1); + + let mut all_tokens = vec![]; + + let mut logits_processor = { + let temperature = args.temperature; + let sampling = if temperature <= 0. { + Sampling::ArgMax + } else { + match (args.top_k, args.top_p) { + (None, None) => Sampling::All { temperature }, + (Some(k), None) => Sampling::TopK { k, temperature }, + (None, Some(p)) => Sampling::TopP { p, temperature }, + (Some(k), Some(p)) => Sampling::TopKThenTopP { k, p, temperature }, + } + }; + LogitsProcessor::from_sampling(args.seed, sampling) + }; + + let start_prompt_processing = std::time::Instant::now(); + + let mut next_token = if !args.split_prompt { + let input = Tensor::new(tokens, &device)?.unsqueeze(0)?; + let logits = model.forward(&input, 0)?; + let logits = logits.squeeze(0)?; + logits_processor.sample(&logits)? + } else { + let mut next_token = 0; + for (pos, token) in tokens.iter().enumerate() { + let input = Tensor::new(&[*token], &device)?.unsqueeze(0)?; + let logits = model.forward(&input, pos)?; + let logits = logits.squeeze(0)?; + next_token = logits_processor.sample(&logits)? + } + next_token + }; + + let prompt_dt = start_prompt_processing.elapsed(); + + all_tokens.push(next_token); + + if let Some(t) = tos.next_token(next_token)? { + print!("{t}"); + std::io::stdout().flush()?; + } + + let eos_token = *tos.tokenizer().get_vocab(true).get("<|im_end|>").unwrap(); + + let start_post_prompt = std::time::Instant::now(); + + let mut sampled = 0; + for index in 0..to_sample { + let input = Tensor::new(&[next_token], &device)?.unsqueeze(0)?; + let logits = model.forward(&input, tokens.len() + index)?; + let logits = logits.squeeze(0)?; + let logits = if args.repeat_penalty == 1. { + logits + } else { + let start_at = all_tokens.len().saturating_sub(args.repeat_last_n); + candle_transformers::utils::apply_repeat_penalty( + &logits, + args.repeat_penalty, + &all_tokens[start_at..], + )? + }; + next_token = logits_processor.sample(&logits)?; + all_tokens.push(next_token); + if let Some(t) = tos.next_token(next_token)? { + print!("{t}"); + std::io::stdout().flush()?; + } + sampled += 1; + if next_token == eos_token { + break; + }; + } + + if let Some(rest) = tos.decode_rest().map_err(candle::Error::msg)? { + print!("{rest}"); + } + + std::io::stdout().flush()?; + let dt = start_post_prompt.elapsed(); + println!( + "\n\n{:4} prompt tokens processed: {:.2} token/s", + tokens.len(), + tokens.len() as f64 / prompt_dt.as_secs_f64(), + ); + println!( + "{sampled:4} tokens generated: {:.2} token/s", + sampled as f64 / dt.as_secs_f64(), + ); + Ok(()) +} diff --git a/candle-transformers/src/models/mod.rs b/candle-transformers/src/models/mod.rs index 27f3b963..790ad439 100644 --- a/candle-transformers/src/models/mod.rs +++ b/candle-transformers/src/models/mod.rs @@ -90,6 +90,7 @@ pub mod quantized_mpt; pub mod quantized_phi; pub mod quantized_phi3; pub mod quantized_qwen2; +pub mod quantized_qwen3; pub mod quantized_recurrent_gemma; pub mod quantized_rwkv_v5; pub mod quantized_rwkv_v6; diff --git a/candle-transformers/src/models/quantized_qwen3.rs b/candle-transformers/src/models/quantized_qwen3.rs new file mode 100644 index 00000000..34dba8cd --- /dev/null +++ b/candle-transformers/src/models/quantized_qwen3.rs @@ -0,0 +1,428 @@ +//! Qwen3 implementation with quantization support. +//! +//! Based on the Qwen3 architecture and implemented with quantized weights +//! for reduced memory usage and faster inference on compatible hardware. +//! +//! References: +//! - [Qwen3 Models](https://huggingface.co/Qwen/Qwen3-0.6B) (architecture based on official implementations) +//! +use super::with_tracing::QMatMul; +use crate::{quantized_nn::RmsNorm, utils::repeat_kv}; +use candle::quantized::{gguf_file, QTensor}; +use candle::{DType, Device, Result, Tensor}; +use candle_nn::{kv_cache::KvCache, Activation, Embedding, Module}; +use std::io::{Read, Seek}; +use std::sync::Arc; + +struct Gguf { + ct: gguf_file::Content, + reader: R, + device: Device, +} + +impl Gguf { + fn new(ct: gguf_file::Content, reader: R, device: Device) -> Self { + Self { ct, reader, device } + } + + fn qmatmul(&mut self, name: &str) -> Result { + let ws = self.ct.tensor(&mut self.reader, name, &self.device)?; + QMatMul::from_weights(ws.into()) + } + + fn rms_norm(&mut self, name: &str, eps: f64) -> Result { + let ws = self.ct.tensor(&mut self.reader, name, &self.device)?; + RmsNorm::from_qtensor(ws, eps) + } + + fn metadata(&self) -> &std::collections::HashMap { + &self.ct.metadata + } + + fn tensor(&mut self, name: &str) -> Result { + self.ct.tensor(&mut self.reader, name, &self.device) + } +} + +#[derive(Debug, Clone)] +struct MlpWeights { + gate_proj: QMatMul, + up_proj: QMatMul, + down_proj: QMatMul, + act_fn: Activation, + span: tracing::Span, +} + +impl MlpWeights { + fn new(gg: &mut Gguf, prefix: &str) -> Result { + let gate_proj = gg.qmatmul(&format!("{prefix}.ffn_gate.weight"))?; + let up_proj = gg.qmatmul(&format!("{prefix}.ffn_up.weight"))?; + let down_proj = gg.qmatmul(&format!("{prefix}.ffn_down.weight"))?; + let act_fn = Activation::Silu; + let span = tracing::span!(tracing::Level::TRACE, "mlp"); + Ok(Self { + gate_proj, + up_proj, + down_proj, + act_fn, + span, + }) + } +} + +impl Module for MlpWeights { + fn forward(&self, x: &Tensor) -> Result { + let _enter = self.span.enter(); + let gate = self.gate_proj.forward(x)?.apply(&self.act_fn)?; + let up = self.up_proj.forward(x)?; + let gated = (gate * up)?; + self.down_proj.forward(&gated) + } +} + +#[derive(Debug, Clone)] +struct RotaryEmbedding { + sin: Tensor, + cos: Tensor, +} + +impl RotaryEmbedding { + fn new( + dtype: DType, + head_dim: usize, + max_position_embeddings: usize, + rope_theta: f64, + dev: &Device, + ) -> Result { + let dim = head_dim; + let max_seq_len = max_position_embeddings; + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / rope_theta.powf(i as f64 / dim as f64) as f32) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(dtype)?; + let t = Tensor::arange(0u32, max_seq_len as u32, dev)? + .to_dtype(dtype)? + .reshape((max_seq_len, 1))?; + let freqs = t.matmul(&inv_freq)?; + Ok(Self { + sin: freqs.sin()?, + cos: freqs.cos()?, + }) + } + + /// Apply RoPE (q, k shape: B x H x L x D) + fn apply(&self, q: &Tensor, k: &Tensor, offset: usize) -> Result<(Tensor, Tensor)> { + let (_, _, seq_len, _) = q.dims4()?; + let cos = self.cos.narrow(0, offset, seq_len)?.to_dtype(q.dtype())?; + let sin = self.sin.narrow(0, offset, seq_len)?.to_dtype(q.dtype())?; + let q_embed = candle_nn::rotary_emb::rope(&q.contiguous()?, &cos, &sin)?; + let k_embed = candle_nn::rotary_emb::rope(&k.contiguous()?, &cos, &sin)?; + Ok((q_embed, k_embed)) + } +} + +#[derive(Debug, Clone)] +struct AttentionWeights { + q_proj: QMatMul, + k_proj: QMatMul, + v_proj: QMatMul, + o_proj: QMatMul, + q_norm: RmsNorm, + k_norm: RmsNorm, + num_heads: usize, + num_kv_heads: usize, + num_kv_groups: usize, + head_dim: usize, + rotary_emb: Arc, + kv_cache: KvCache, + span_attn: tracing::Span, +} + +impl AttentionWeights { + fn new( + gg: &mut Gguf, + num_heads: usize, + num_kv_heads: usize, + head_dim: usize, + rms_norm_eps: f64, + rotary_emb: Arc, + prefix: &str, + ) -> Result { + let num_kv_groups = num_heads / num_kv_heads; + + let q_proj = gg.qmatmul(&format!("{prefix}.attn_q.weight"))?; + let k_proj = gg.qmatmul(&format!("{prefix}.attn_k.weight"))?; + let v_proj = gg.qmatmul(&format!("{prefix}.attn_v.weight"))?; + let o_proj = gg.qmatmul(&format!("{prefix}.attn_output.weight"))?; + + let q_norm = gg.rms_norm(&format!("{prefix}.attn_q_norm.weight"), rms_norm_eps)?; + let k_norm = gg.rms_norm(&format!("{prefix}.attn_k_norm.weight"), rms_norm_eps)?; + + let max_position_embeddings = gg + .metadata() + .get("qwen3.context_length") + .and_then(|v| v.to_u32().ok()) + .unwrap_or(4096) as usize; + let kv_cache = KvCache::new(2, max_position_embeddings); + + let span_attn = tracing::span!(tracing::Level::TRACE, "attn"); + + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + q_norm, + k_norm, + num_heads, + num_kv_heads, + num_kv_groups, + head_dim, + rotary_emb, + kv_cache, + span_attn, + }) + } + + fn forward(&mut self, x: &Tensor, attn_mask: Option<&Tensor>, offset: usize) -> Result { + let _enter = self.span_attn.enter(); + let (b, l, _) = x.dims3()?; + + let q = self.q_proj.forward(x)?; + let k = self.k_proj.forward(x)?; + let v = self.v_proj.forward(x)?; + + let q = q + .reshape((b, l, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + let k = k + .reshape((b, l, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let v = v + .reshape((b, l, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + + let q_flat = q.flatten(0, 2)?; + let k_flat = k.flatten(0, 2)?; + + let q_flat = self.q_norm.forward(&q_flat)?; + let k_flat = self.k_norm.forward(&k_flat)?; + let q = q_flat.reshape((b, self.num_heads, l, self.head_dim))?; + let k = k_flat.reshape((b, self.num_kv_heads, l, self.head_dim))?; + + let (q, k) = self.rotary_emb.apply(&q, &k, offset)?; + + // Reset KV cache if we're at the first position + if offset == 0 { + self.kv_cache.reset(); + } + let (k, v) = self.kv_cache.append(&k.contiguous()?, &v.contiguous()?)?; + + let k = repeat_kv(k, self.num_kv_groups)?.contiguous()?; + let v = repeat_kv(v, self.num_kv_groups)?.contiguous()?; + + let scale = 1.0 / (self.head_dim as f64).sqrt(); + let mut scores = (q.matmul(&k.transpose(2, 3)?)? * scale)?; + if let Some(m) = attn_mask { + let m_dtype = m.dtype(); + let scores_dtype = scores.dtype(); + let mask = if m_dtype != scores_dtype { + m.to_dtype(scores_dtype)? + } else { + m.clone() + }; + scores = scores.broadcast_add(&mask)?; + } + let probs = candle_nn::ops::softmax_last_dim(&scores)?; + let ctx = probs.matmul(&v)?; // (B, H, L, D) + let reshaped_ctx = ctx + .transpose(1, 2)? + .reshape((b, l, self.num_heads * self.head_dim))?; + self.o_proj.forward(&reshaped_ctx) + } +} + +#[derive(Debug, Clone)] +struct LayerWeights { + self_attn: AttentionWeights, + mlp: MlpWeights, + ln1: RmsNorm, + ln2: RmsNorm, +} + +impl LayerWeights { + fn new( + gg: &mut Gguf, + num_attention_heads: usize, + num_key_value_heads: usize, + head_dim: usize, + rms_norm_eps: f64, + rotary: Arc, + layer_idx: usize, + ) -> Result { + let prefix = format!("blk.{layer_idx}"); + + let ln1 = gg.rms_norm(&format!("{prefix}.attn_norm.weight"), rms_norm_eps)?; + let ln2 = gg.rms_norm(&format!("{prefix}.ffn_norm.weight"), rms_norm_eps)?; + let self_attn = AttentionWeights::new( + gg, + num_attention_heads, + num_key_value_heads, + head_dim, + rms_norm_eps, + rotary, + &prefix, + )?; + let mlp = MlpWeights::new(gg, &prefix)?; + Ok(Self { + self_attn, + mlp, + ln1, + ln2, + }) + } + + fn forward(&mut self, x: &Tensor, mask: Option<&Tensor>, offset: usize) -> Result { + let h = self.ln1.forward(x)?; + let h = self.self_attn.forward(&h, mask, offset)?; + let x = (x + h)?; + let h2 = self.ln2.forward(&x)?; + let h2 = h2.apply(&self.mlp)?; + x + h2 + } +} + +#[derive(Debug, Clone)] +pub struct ModelWeights { + embed_tokens: Embedding, + layers: Vec, + norm: RmsNorm, + lm_head: QMatMul, + device: Device, + dtype: DType, + span: tracing::Span, + span_output: tracing::Span, +} + +impl ModelWeights { + pub fn from_gguf( + ct: gguf_file::Content, + reader: &mut R, + device: &Device, + ) -> Result { + let mut gg = Gguf::new(ct, reader, device.clone()); + let md_get = |s: &str| match gg.metadata().get(s) { + None => candle::bail!("cannot find {s} in metadata"), + Some(v) => Ok(v), + }; + + let num_attention_heads = md_get("qwen3.attention.head_count")?.to_u32()? as usize; + let num_kv_heads = md_get("qwen3.attention.head_count_kv")?.to_u32()? as usize; + let head_dim = md_get("qwen3.attention.key_length")?.to_u32()? as usize; + let num_layers = md_get("qwen3.block_count")?.to_u32()? as usize; + let hidden_size = md_get("qwen3.embedding_length")?.to_u32()? as usize; + let max_position_embeddings = md_get("qwen3.context_length")?.to_u32()? as usize; + let rms_norm_eps = md_get("qwen3.attention.layer_norm_rms_epsilon")?.to_f32()? as f64; + let rope_freq_base = md_get("qwen3.rope.freq_base")?.to_f32()? as f64; + + let dtype = match gg.metadata().get("general.dtype") { + Some(v) => match v.to_u32() { + Ok(0) => DType::F32, + Ok(1) => DType::F16, + _ => DType::F16, + }, + None => DType::F16, + }; + + let embed_tensor = gg.tensor("token_embd.weight")?; + let embed_tokens = Embedding::new(embed_tensor.dequantize(device)?, hidden_size); + + let rotary = Arc::new(RotaryEmbedding::new( + dtype, + head_dim, + max_position_embeddings, + rope_freq_base, + device, + )?); + + let mut layers = Vec::with_capacity(num_layers); + for i in 0..num_layers { + layers.push(LayerWeights::new( + &mut gg, + num_attention_heads, + num_kv_heads, + head_dim, + rms_norm_eps, + rotary.clone(), + i, + )?); + } + + let norm = gg.rms_norm("output_norm.weight", rms_norm_eps)?; + // Load output projection tensor, falling back to tied embeddings like gemma3 + let lm_head_tensor = match gg.tensor("output.weight") { + Ok(tensor) => tensor, + Err(_) => gg.tensor("token_embd.weight")?, + }; + let lm_head = QMatMul::from_weights(lm_head_tensor.into())?; + let span = tracing::span!(tracing::Level::TRACE, "model"); + let span_output = tracing::span!(tracing::Level::TRACE, "output"); + Ok(Self { + embed_tokens, + layers, + norm, + lm_head, + device: device.clone(), + dtype, + span, + span_output, + }) + } + + fn causal_mask( + &self, + b: usize, + tgt: usize, + offset: usize, + sw: Option, + ) -> Result { + let minf = f32::NEG_INFINITY; + let mask: Vec<_> = (0..tgt) + .flat_map(|i| { + (0..(tgt + offset)).map(move |j| { + let past_ok = j <= i + offset; + let sw_ok = match sw { + Some(w) => (i + offset) as i64 - j as i64 <= w as i64, + None => true, + }; + if past_ok && sw_ok { + 0. + } else { + minf + } + }) + }) + .collect(); + Tensor::from_slice(&mask, (b, 1, tgt, tgt + offset), &self.device)?.to_dtype(self.dtype) + } + + pub fn forward(&mut self, input: &Tensor, offset: usize) -> Result { + let _enter = self.span.enter(); + let (b, l) = input.dims2()?; + let mut h = self.embed_tokens.forward(input)?; + let causal_mask = if l == 1 { + None + } else { + Some(self.causal_mask(b, l, offset, None)?) + }; + for layer in &mut self.layers { + h = layer.forward(&h, causal_mask.as_ref(), offset)?; + } + let h = self.norm.forward(&h)?; + let _enter = self.span_output.enter(); + let last_hidden = h.narrow(1, l - 1, 1)?; + self.lm_head.forward(&last_hidden)?.squeeze(1) + } +} diff --git a/candle-transformers/src/models/qwen3.rs b/candle-transformers/src/models/qwen3.rs index 30ea3c15..dd90b193 100644 --- a/candle-transformers/src/models/qwen3.rs +++ b/candle-transformers/src/models/qwen3.rs @@ -53,7 +53,7 @@ impl Qwen3RotaryEmbedding { } /// Apply RoPE (q, k shape: B x H x L x D) - fn apply(&self, q: &Tensor, k: &Tensor, offset: usize) -> Result<(Tensor, Tensor)> { + pub(crate) fn apply(&self, q: &Tensor, k: &Tensor, offset: usize) -> Result<(Tensor, Tensor)> { let (_, _, seq_len, _) = q.dims4()?; let cos = self.cos.narrow(0, offset, seq_len)?; let sin = self.sin.narrow(0, offset, seq_len)?;