Repository: huggingface/candle
File: candle-transformers/src/models/qwen3.rs (Rust)

Code before:
impl Module for Qwen3MLP {
    fn forward(&self, x: &Tensor) -> Result<Tensor> {
        let lhs = x.apply(&self.gate_proj)?.apply(&self.act_fn)?;
        let rhs = x.apply(&self.up_proj)?;
        (lhs * rhs)?.apply(&self.down_proj)
    }
}

#[derive(Debug, Clone)]
struct Qwen3Attention {
    // projections
    q_proj: Linear,
    k_proj: Linear,
    v_proj: Linear,
    o_proj: Linear,
    // norms
    q_norm: Qwen3HeadRmsNorm,
    k_norm: Qwen3HeadRmsNorm,
    // hyper params
    num_heads: usize,
    num_kv_heads: usize,
    num_kv_groups: usize,
    head_dim: usize,
    hidden_size: usize,
    // sliding window
    _sliding_window: Option<usize>,
    // utils
    rotary_emb: Arc<Qwen3RotaryEmbedding>,
    kv_cache: Option<(Tensor, Tensor)>,
}

impl Qwen3Attention {
    fn new(
        cfg: &Config,
        rotary_emb: Arc<Qwen3RotaryEmbedding>,
        layer_idx: usize,
        vb: VarBuilder,
    ) -> Result<Self> {
        let head_dim = cfg.head_dim;
        let num_heads = cfg.num_attention_heads;
        let num_kv_heads = cfg.num_key_value_heads;
        let num_kv_groups = num_heads / num_kv_heads;

        let (q_proj, k_proj, v_proj, o_proj) = if cfg.attention_bias {
            (
                linear(cfg.hidden_size, num_heads * head_dim, vb.pp("q_proj"))?,
                linear(cfg.hidden_size, num_kv_heads * head_dim, vb.pp("k_proj"))?,
                linear(cfg.hidden_size, num_kv_heads * head_dim, vb.pp("v_proj"))?,
                linear(num_heads * head_dim, cfg.hidden_size, vb.pp("o_proj"))?,
            )
        } else {

Review comment (suggestion):
Why is this prefixed with an `_`, is it unused at the moment, if so a comment should probably explain why it doesn't matter.

Code after:
            vb.pp("k_proj"),
        )?;
        let v_proj = linear_b(
            cfg.hidden_size,
            num_kv_heads * head_dim,
            cfg.attention_bias,
            vb.pp("v_proj"),
        )?;
        let o_proj = linear_b(
            num_heads * head_dim,
            cfg.hidden_size,
            cfg.attention_bias,
            vb.pp("o_proj"),
        )?;

        let q_norm = RmsNorm::new(head_dim, cfg.rms_norm_eps, vb.pp("q_norm"))?;
        let k_norm = RmsNorm::new(head_dim, cfg.rms_norm_eps, vb.pp("k_norm"))?;

        // Necessary because the hidden_size in the config isn't always accurate
        let hidden_size = head_dim * cfg.num_attention_heads;

        let kv_cache = KvCache::new(2, cfg.max_position_embeddings);

        Ok(Self {
            q_proj,
            k_proj,
            v_proj,
            o_proj,
            q_norm,
            k_norm,
            num_heads,
            num_kv_heads,
            num_kv_groups,
            head_dim,
            hidden_size,
            rotary_emb,
            kv_cache,
        })
    }

    pub(crate) fn forward(
        &mut self,
        x: &Tensor,
        attn_mask: Option<&Tensor>,
        offset: usize,
    ) -> Result<Tensor> {
        let (b, l, _) = x.dims3()?;

        // 1. Proj
        let q = self.q_proj.forward(x)?;
        let k = self.k_proj.forward(x)?;