ronantakizawa/codereview-bench huggingface dataset:
Repository: hpcaitech/ColossalAI
File: extensions/csrc/cuda/get_cos_and_sin_kernel.cu (CUDA)

Code before:
#include <ATen/cuda/CUDAContext.h>
#include <torch/extension.h>

#include "utils/vector_copy_utils.h"
#include "../common/micros.h"
#include "stdio.h"

template <typename scalar_t, bool Aligned, int VecSize>
__global__ void apply_get_context_cos_and_sin_kernel(
    scalar_t* __restrict__ cos,
    scalar_t* __restrict__ sin,
    const scalar_t* __restrict__ cos_cache_ptr,
    const scalar_t* __restrict__ sin_cache_ptr,
    const int* __restrict__ sequence_lengths,
    const int* __restrict__ comsum_lengths,
    const int batch_size,
    const int head_dim
) {
    int token_id = blockIdx.x;
    if ( token_id >= sequence_lengths[blockIdx.y] ) {
        return ;
    }

    int begin_id = threadIdx.x * VecSize;
    int src_offset_id = token_id * head_dim;
    int dest_offset_id = src_offset_id;

    if (blockIdx.y > 0) {
        dest_offset_id += comsum_lengths[blockIdx.y - 1] * head_dim;
    }

    for (; begin_id <= head_dim - VecSize; begin_id += blockDim.x){
        copy_vector<scalar_t, VecSize>(cos + dest_offset_id + begin_id, cos_cache_ptr + src_offset_id + begin_id);

Review comment (suggestion):
Just wonder why not putting `Aligned` and `VecSize` as function parameters ,instead of template parameters? Using a non-manifestly-constant-evaluated expression as a non-type template parameter is not recommended. This is actually causing potential issues we discussed before in this [PR](https://github.com/hpcaitech/ColossalAI/pull/5461). 

Code after:
#include <ATen/cuda/CUDAContext.h>
#include <torch/extension.h>

#include "utils/vector_copy_utils.h"
#include "../common/micros.h"
#include "stdio.h"

template <typename scalar_t, bool Aligned, int VecSize>
__device__ void apply_cos_and_sin_memcopy(
    scalar_t* __restrict__ cos,
    scalar_t* __restrict__ sin,
    const scalar_t* __restrict__ cos_cache_ptr,
    const scalar_t* __restrict__ sin_cache_ptr,
    const int* __restrict__ sequence_lengths,
    const int head_dim,
    const int dest_offset_id,
    const int src_offset_id
 ) {

    int begin_id = threadIdx.x * VecSize;

    for (; begin_id <= head_dim - VecSize; begin_id += blockDim.x){
        copy_vector<scalar_t, VecSize>(cos + dest_offset_id + begin_id, cos_cache_ptr + src_offset_id + begin_id);
        copy_vector<scalar_t, VecSize>(sin + dest_offset_id + begin_id, sin_cache_ptr + src_offset_id + begin_id);
    }

    if (!Aligned) {
        for (; begin_id < head_dim; ++begin_id ) {
            cos[dest_offset_id + begin_id] = cos_cache_ptr[src_offset_id + begin_id];
            sin[dest_offset_id + begin_id] = sin_cache_ptr[src_offset_id + begin_id];
        }
    }
}
Repository: huggingface/candle
File: candle-transformers/src/models/whisper/audio.rs (Rust)

Code before:
                &samples_clone,
                &filters_clone,
                fft_size,
                fft_step,
                speed_up,
                n_len,
                n_mel,
                n_threads,
            )
        }));
    }

    let mut all_outputs = Vec::new();
    for handle in handles {
        match handle.join() {
            Ok(output) => all_outputs.push(output),
            Err(_) => panic!("Thread failed"),
        }
    }

    let l = all_outputs[0].len();
    let mut mel = vec![zero; l];

    for i in (0..l).step_by(n_threads) {
        for ouputs in all_outputs.iter() {
            for j in 0..n_threads {
                mel[i + j] += ouputs[i + j];
            }
        }
    }

    let mmax = mel
        .iter()
        .max_by(|&u, &v| u.partial_cmp(v).unwrap_or(std::cmp::Ordering::Greater))
        .copied()
        .unwrap_or(zero)
        - T::from(8).unwrap();
    for m in mel.iter_mut() {
        let v = T::max(*m, mmax);
        *m = v / four + one
    }
    mel
}

pub fn pcm_to_mel<T: Float + std::fmt::Display + Send + Sync + 'static>(
    cfg: &super::Config,
    samples: &[T],
    filters: &[T],
) -> Vec<T> {
    log_mel_spectrogram_(
        samples,

Review comment (suggestion):
Could you add some comment to detail a bit what is going on here, the first loop on the outputs of all threads then iterating on the number of thtreads is pretty hard to follow.

Code after:
            .collect::<Vec<_>>()
    });

    let l = all_outputs[0].len();
    let mut mel = vec![zero; l];

    // iterate over mel spectrogram segments, dividing work by threads.
    for segment_start in (0..l).step_by(n_threads) {
        // go through each thread's output.
        for thread_output in all_outputs.iter() {
            // add each thread's piece to our mel spectrogram.
            for offset in 0..n_threads {
                let mel_index = segment_start + offset; // find location in mel.
                if mel_index < mel.len() {
                    // Make sure we don't go out of bounds.
                    mel[mel_index] += thread_output[mel_index];
                }
            }
        }
    }

    let mmax = mel
        .iter()
        .max_by(|&u, &v| u.partial_cmp(v).unwrap_or(std::cmp::Ordering::Greater))
        .copied()
        .unwrap_or(zero)
        - T::from(8).unwrap();
    for m in mel.iter_mut() {
        let v = T::max(*m, mmax);
        *m = v / four + one
    }
    mel
}

pub fn pcm_to_mel<T: Float>(cfg: &super::Config, samples: &[T], filters: &[T]) -> Vec<T> {
    log_mel_spectrogram_(
        samples,
        filters,
        super::N_FFT,
        super::HOP_LENGTH,
        cfg.num_mel_bins,
        false,
    )
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_fft() {
Repository: huggingface/candle
File: candle-transformers/src/models/quantized_mixformer.rs (Rust)

Code before:
        let _enter = self.span.enter();
        let (_b_size, seq_len) = xs.dims2()?;
        let mut xs = xs.apply(&self.embedding)?;
        let mask = if seq_len <= 1 {
            None
        } else {
            Some(get_mask(seq_len, xs.device())?)
        };
        for block in self.blocks.iter_mut() {
            xs = block.forward(&xs, mask.as_ref())?;
        }
        xs.narrow(1, seq_len - 1, 1)?.apply(&self.head)?.squeeze(1)
    }

    pub fn forward_with_img(
        &mut self,
        bos_token: &Tensor,
        xs: &Tensor,
        img_embeds: &Tensor,
    ) -> Result<Tensor> {
        let _enter = self.span.enter();
        let xs = xs.apply(&self.embedding)?;
        let bos_token = bos_token.apply(&self.embedding)?;
        let mut xs = Tensor::cat(&[bos_token, img_embeds.clone(), xs], 1)?;
        let (_b_size, seq_len, _embds) = xs.dims3()?;
        let mask = if seq_len <= 1 {
            None
        } else {
            Some(get_mask(seq_len, xs.device())?)
        };
        for block in self.blocks.iter_mut() {
            xs = block.forward(&xs, mask.as_ref())?
        }
        let xs = xs
            .narrow(1, seq_len - 1, 1)?
            .apply(&self.head)?
            .squeeze(1)?;
        Ok(xs)
    }

    pub fn clear_kv_cache(&mut self) {
        self.blocks.iter_mut().for_each(|b| b.clear_kv_cache())
    }
}


Review comment (suggestion):
I don't think `seq_len` can ever be <= 1 here, so let's remove that bit as well as the equivalent part in `mixformers.rs`.

Code after:
        let _enter = self.span.enter();
        let (_b_size, seq_len) = xs.dims2()?;
        let mut xs = xs.apply(&self.embedding)?;
        let mask = if seq_len <= 1 {
            None
        } else {
            Some(get_mask(seq_len, xs.device())?)
        };
        for block in self.blocks.iter_mut() {
            xs = block.forward(&xs, mask.as_ref())?;
        }
        xs.narrow(1, seq_len - 1, 1)?.apply(&self.head)?.squeeze(1)
    }

    pub fn forward_with_img(
        &mut self,
        bos_token: &Tensor,
        xs: &Tensor,
        img_embeds: &Tensor,
    ) -> Result<Tensor> {
        let _enter = self.span.enter();
        let xs = xs.apply(&self.embedding)?;
        let bos_token = bos_token.apply(&self.embedding)?;
        // Python implementation sequence order is <bos token embedding><img embedding><rest of text embedding>
        // https://github.com/vikhyat/moondream/blob/a9d788a20d1543fb1479edc54106e88cff7759d3/moondream/moondream.py#L43-L56
        let mut xs = Tensor::cat(&[bos_token, img_embeds.clone(), xs], 1)?;
        let (_b_size, seq_len, _embds) = xs.dims3()?;
        let mask = Some(get_mask(seq_len, xs.device())?);
        for block in self.blocks.iter_mut() {
            xs = block.forward(&xs, mask.as_ref())?
        }
        let xs = xs
            .narrow(1, seq_len - 1, 1)?
            .apply(&self.head)?
            .squeeze(1)?;
        Ok(xs)
    }

    pub fn clear_kv_cache(&mut self) {
        self.blocks.iter_mut().for_each(|b| b.clear_kv_cache())
    }
}
Repository: huggingface/candle
File: candle-metal-kernels/src/lib.rs (Rust)

Code before:
    encoder.use_resource(rhs_buffer, metal::MTLResourceUsage::Read);
    encoder.use_resource(output, metal::MTLResourceUsage::Write);
    encoder.dispatch_thread_groups(grid_size, group_size);
    Ok(())
}

pub fn call_const_fill(
    device: &Device,
    ep: impl EncoderProvider,
    kernels: &Kernels,
    name: &'static str,
    length: usize,
    output: &Buffer,
    v: f32
) -> Result<(), MetalKernelError> {
    let pipeline = kernels.load_pipeline(device, Source::Fill, name)?;
    let encoder = ep.encoder();
    let encoder: &ComputeCommandEncoderRef = encoder.as_ref();

    encoder.set_compute_pipeline_state(&pipeline);

    set_params!(encoder, (output, v));

    encoder.use_resource(output, metal::MTLResourceUsage::Write);

    let grid_size = MTLSize { width: length as u64, height: 1, depth: 1 };
    let thread_group_size = MTLSize { width: pipeline.max_total_threads_per_threadgroup(), height: 1, depth: 1 };

    encoder.dispatch_threads(grid_size, thread_group_size);

    Ok(())
}

#[cfg(test)]
mod tests;


Review comment (suggestion):
Is `length` here correct? Shouldn't it divided by the number of threads? Anyway better to use `utils::linear_split` as it should do the proper thing here.

Code after:
    encoder.use_resource(rhs_buffer, metal::MTLResourceUsage::Read);
    encoder.use_resource(output, metal::MTLResourceUsage::Write);
    encoder.dispatch_thread_groups(grid_size, group_size);
    Ok(())
}

pub fn call_const_fill(
    device: &Device,
    ep: impl EncoderProvider,
    kernels: &Kernels,
    name: &'static str,
    length: usize,
    output: &Buffer,
    v: f32,
) -> Result<(), MetalKernelError> {
    let pipeline = kernels.load_pipeline(device, Source::Fill, name)?;
    let encoder = ep.encoder();
    let encoder: &ComputeCommandEncoderRef = encoder.as_ref();

    encoder.set_compute_pipeline_state(&pipeline);

    set_params!(encoder, (output, v, length));

    let (thread_group_count, thread_group_size) = linear_split(&pipeline, length);

    encoder.use_resource(output, metal::MTLResourceUsage::Write);
    encoder.dispatch_thread_groups(thread_group_count, thread_group_size);

    Ok(())
}

#[cfg(test)]
mod tests;
Repository: huggingface/candle
File: candle-transformers/src/models/mmdit/model.rs (Rust)

Code before:
            patch_size: 2,
            in_channels: 16,
            out_channels: 16,
            depth: 24,
            head_size: 64,
            adm_in_channels: 2048,
            pos_embed_max_size: 192,
            context_embed_size: 4096,
            frequency_embedding_size: 256,
        }
    }

    pub fn sd3_5_large() -> Self {
        Self {
            patch_size: 2,
            in_channels: 16,
            out_channels: 16,
            depth: 38,
            head_size: 64,
            adm_in_channels: 2048,
            pos_embed_max_size: 192,
            context_embed_size: 4096,
            frequency_embedding_size: 256,
        }
    }
}

pub struct MMDiT {
    core: MMDiTCore,
    patch_embedder: PatchEmbedder,
    pos_embedder: PositionEmbedder,
    timestep_embedder: TimestepEmbedder,
    vector_embedder: VectorEmbedder,
    context_embedder: nn::Linear,
    unpatchifier: Unpatchifier,
}

impl MMDiT {
    pub fn new(cfg: &Config, use_flash_attn: bool, vb: nn::VarBuilder) -> Result<Self> {
        let hidden_size = cfg.head_size * cfg.depth;
        let core = MMDiTCore::new(
            cfg.depth,
            hidden_size,
            cfg.depth,
            cfg.patch_size,
            cfg.out_channels,
            use_flash_attn,
            vb.clone(),
        )?;
        let patch_embedder = PatchEmbedder::new(
            cfg.patch_size,

Review comment (suggestion):
Notably, `pos_embed_max_size` for position embedding has been increased for SD3.5-medium, making it even bigger than SD3.5-large, which kept the original size of that of SD3-medium...

Code after:
}

impl Config {
    pub fn sd3_medium() -> Self {
        Self {
            patch_size: 2,
            in_channels: 16,
            out_channels: 16,
            depth: 24,
            head_size: 64,
            adm_in_channels: 2048,
            pos_embed_max_size: 192,
            context_embed_size: 4096,
            frequency_embedding_size: 256,
        }
    }

    pub fn sd3_5_medium() -> Self {
        Self {
            patch_size: 2,
            in_channels: 16,
            out_channels: 16,
            depth: 24,
            head_size: 64,
            adm_in_channels: 2048,
            pos_embed_max_size: 384,
            context_embed_size: 4096,
            frequency_embedding_size: 256,
        }
    }

    pub fn sd3_5_large() -> Self {
        Self {
            patch_size: 2,
            in_channels: 16,
            out_channels: 16,
            depth: 38,
            head_size: 64,
            adm_in_channels: 2048,
            pos_embed_max_size: 192,
            context_embed_size: 4096,
            frequency_embedding_size: 256,
        }
    }
}

pub struct MMDiT {
    core: MMDiTCore,
    patch_embedder: PatchEmbedder,
    pos_embedder: PositionEmbedder,
    timestep_embedder: TimestepEmbedder,
Repository: huggingface/candle
File: candle-transformers/src/models/qwen3.rs (Rust)

Code before:
        Ok((xs.broadcast_mul(&ws))?.to_dtype(orig_dtype)?)
    }
}

#[derive(Debug, Clone)]
struct Qwen3HeadRmsNorm {
    weight: Tensor,
    eps: f64,
}

impl Qwen3HeadRmsNorm {
    fn new(dim: usize, eps: f64, vb: VarBuilder) -> Result<Self> {
        let weight = vb.get((dim,), "weight")?;
        Ok(Self { weight, eps })
    }

    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
        // xs: (B*L*H, D)
        let orig_dtype = xs.dtype();
        let xs = xs.to_dtype(DType::F32)?;
        let var = (xs.clone() * &xs)?.mean_keepdim(D::Minus1)?;
        let eps_tensor = Tensor::new::<f32>(self.eps as f32, xs.device())?;
        let rms = var.broadcast_add(&eps_tensor)?.powf(-0.5)?;
        let xs = xs.broadcast_mul(&rms)?;
        let ws = self.weight.unsqueeze(0)?;
        Ok((xs.broadcast_mul(&ws))?.to_dtype(orig_dtype)?)
    }
}

fn repeat_kv(kv: &Tensor, n_rep: usize) -> Result<Tensor> {
    if n_rep == 1 {
        return Ok(kv.clone());
    }
    let (b, h_kv, l, d) = kv.dims4()?;
    kv.unsqueeze(2)?
        .expand((b, h_kv, n_rep, l, d))?
        .reshape((b, h_kv * n_rep, l, d))
}

#[derive(Debug, Clone)]
struct Qwen3MLP {
    gate_proj: Linear,
    up_proj: Linear,
    down_proj: Linear,
    act_fn: Activation,
}

impl Qwen3MLP {
    fn new(cfg: &Config, vb: VarBuilder) -> Result<Self> {
        Ok(Self {
            gate_proj: linear_no_bias(cfg.hidden_size, cfg.intermediate_size, vb.pp("gate_proj"))?,

Review comment (suggestion):
```suggestion
        xs.broadcast_mul(&ws)?.to_dtype(orig_dtype)
```

I think.

Code after:
    }
}

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)]
pub(crate) struct Qwen3Attention {
    // projections
    q_proj: Linear,
    k_proj: Linear,
    v_proj: Linear,
    o_proj: Linear,
    // norms
    q_norm: RmsNorm,
    k_norm: RmsNorm,
    // hyper params
    num_heads: usize,
    num_kv_heads: usize,
    num_kv_groups: usize,
    head_dim: usize,
    hidden_size: usize,
    // utils
    rotary_emb: Arc<Qwen3RotaryEmbedding>,
    kv_cache: KvCache,
}

impl Qwen3Attention {
    pub(crate) fn new(
        cfg: &Config,
        rotary_emb: Arc<Qwen3RotaryEmbedding>,
        vb: VarBuilder,
    ) -> Result<Self> {
        if cfg.use_sliding_window {
            candle::bail!("sliding window is not suppored")
        }

        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 = linear_b(
            cfg.hidden_size,
            num_heads * head_dim,
            cfg.attention_bias,
Repository: huggingface/candle
File: candle-transformers/src/models/qwen3.rs (Rust)

Code before:

#[derive(Debug, Clone)]
struct Qwen3RmsNorm {
    weight: Tensor,
    eps: f64,
}

impl Qwen3RmsNorm {
    fn new(dim: usize, eps: f64, vb: VarBuilder) -> Result<Self> {
        let weight = vb.get((dim,), "weight")?;
        Ok(Self { weight, eps })
    }

    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
        let orig_dtype = xs.dtype();
        let xs = xs.to_dtype(DType::F32)?;
        let var = (xs.clone() * &xs)?.mean_keepdim(D::Minus1)?;
        let rms = (var + self.eps)?.powf(-0.5)?;
        let xs = xs.broadcast_mul(&rms)?;
        let ws = self.weight.unsqueeze(0)?.unsqueeze(1)?;
        xs.broadcast_mul(&ws)?.to_dtype(orig_dtype)
    }
}

#[derive(Debug, Clone)]
struct Qwen3HeadRmsNorm {
    weight: Tensor,
    eps: f64,
}

impl Qwen3HeadRmsNorm {
    fn new(dim: usize, eps: f64, vb: VarBuilder) -> Result<Self> {
        let weight = vb.get((dim,), "weight")?;
        Ok(Self { weight, eps })
    }

    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
        // xs: (B*L*H, D)
        let orig_dtype = xs.dtype();
        let xs = xs.to_dtype(DType::F32)?;
        let var = (xs.clone() * &xs)?.mean_keepdim(D::Minus1)?;
        let rms = (var + self.eps)?.powf(-0.5)?;
        let xs = xs.broadcast_mul(&rms)?;
        let ws = self.weight.unsqueeze(0)?;
        xs.broadcast_mul(&ws)?.to_dtype(orig_dtype)
    }
}

fn repeat_kv(kv: &Tensor, n_rep: usize) -> Result<Tensor> {
    if n_rep == 1 {
        return Ok(kv.clone());

Review comment (performance):
Same question as above, should this be the standard `RmsNorm` module? Or maybe just a wrapper around it? The `forward` below is likely to be much less efficient.

Code after:
        Ok((q_embed, k_embed))
    }
}

#[derive(Debug, Clone)]
pub(crate) struct Qwen3MLP {
    gate_proj: Linear,
    up_proj: Linear,
    down_proj: Linear,
    act_fn: Activation,
}

impl Qwen3MLP {
    pub(crate) fn new(cfg: &Config, vb: VarBuilder) -> Result<Self> {
        Ok(Self {
            gate_proj: linear_no_bias(cfg.hidden_size, cfg.intermediate_size, vb.pp("gate_proj"))?,
            up_proj: linear_no_bias(cfg.hidden_size, cfg.intermediate_size, vb.pp("up_proj"))?,
            down_proj: linear_no_bias(cfg.intermediate_size, cfg.hidden_size, vb.pp("down_proj"))?,
            act_fn: cfg.hidden_act,
        })
    }
}

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)]
pub(crate) struct Qwen3Attention {
    // projections
    q_proj: Linear,
    k_proj: Linear,
    v_proj: Linear,
    o_proj: Linear,
    // norms
    q_norm: RmsNorm,
    k_norm: RmsNorm,
    // hyper params
    num_heads: usize,
    num_kv_heads: usize,
    num_kv_groups: usize,
    head_dim: usize,
    hidden_size: usize,
    // utils
    rotary_emb: Arc<Qwen3RotaryEmbedding>,
    kv_cache: KvCache,
}
Repository: huggingface/candle
File: candle-transformers/src/models/qwen3.rs (Rust)

Code before:
            cfg.attention_bias,
            vb.pp("q_proj"),
        )?;
        let k_proj = linear_b(
            cfg.hidden_size,
            num_kv_heads * head_dim,
            cfg.attention_bias,
            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 cofig 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>,

Review comment (suggestion):
```suggestion
        // Necessary because the hidden_size in the config isn't always accurate
```

Code after:
            cfg.attention_bias,
            vb.pp("q_proj"),
        )?;
        let k_proj = linear_b(
            cfg.hidden_size,
            num_kv_heads * head_dim,
            cfg.attention_bias,
            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>,