Method forward

cake-core/src/models/common/attention.rs:152–357

Process the input tensor using the given state indexes and cache.

(
        &self,
        x: &Tensor,
        index_pos: usize,
        block_idx: usize,
        cache: &mut super::Cache,
    )

    /// Process the input tensor using the given state indexes and cache.
    pub fn forward(
        &self,
        x: &Tensor,
        index_pos: usize,
        block_idx: usize,
        cache: &mut super::Cache,
    ) -> anyhow::Result<Tensor> {
        let (b_sz, seq_len, _hidden_size) = x.dims3().map_err(|e| anyhow!("x.dims3 -> {e}"))?;

        // Single fused QKV projection (routed through backend for GPU acceleration)
        let qkv = self
            .backend.linear_forward(x, &self.qkv_proj_weight, self.qkv_proj_bias.as_ref())
            .map_err(|e| anyhow!("qkv.forward -> {e}"))?;

        let q = qkv
            .narrow(D::Minus1, 0, self.size_q)
            .map_err(|e| anyhow!("q split -> {e}"))?;
        let k = qkv
            .narrow(D::Minus1, self.size_q, self.size_kv)
            .map_err(|e| anyhow!("k split -> {e}"))?;
        let v = qkv
            .narrow(D::Minus1, self.size_q + self.size_kv, self.size_kv)
            .map_err(|e| anyhow!("v split -> {e}"))?;

        // OLMo2-style: apply QK-norm BEFORE head reshape (norm dim = size_q/size_kv).
        let (q, k) = if self.pre_reshape_qk_norm {
            let q = if let Some(w) = &self.q_norm_weight {
                self.backend.rms_norm(&q.contiguous()
                    .map_err(|e| anyhow!("pre_reshape q contiguous -> {e}"))?, w, self.qk_norm_eps)
                    .map_err(|e| anyhow!("pre_reshape q_norm -> {e}"))?
            } else { q };
            let k = if let Some(w) = &self.k_norm_weight {
                self.backend.rms_norm(&k.contiguous()
                    .map_err(|e| anyhow!("pre_reshape k contiguous -> {e}"))?, w, self.qk_norm_eps)
                    .map_err(|e| anyhow!("pre_reshape k_norm -> {e}"))?
            } else { k };
            (q, k)
        } else {
            (q, k)
        };

        // Reshape: (b, seq, heads, head_dim)
        let q = q
            .reshape((b_sz, seq_len, self.num_attention_heads, self.head_dim))?;
        let k = k
            .reshape((b_sz, seq_len, self.num_key_value_heads, self.head_dim))?;
        let v = v
            .reshape((b_sz, seq_len, self.num_key_value_heads, self.head_dim))?;

        // Standard QK-norm: applied after reshape (on head_dim, last dim) before transpose.
        let q = if !self.pre_reshape_qk_norm {
            if let Some(w) = &self.q_norm_weight {
                self.backend.rms_norm(&q.contiguous()
                    .map_err(|e| anyhow!("q contiguous -> {e}"))?, w, self.qk_norm_eps)
                    .map_err(|e| anyhow!("q_norm -> {e}"))?
            } else { q }
        } else { q };
        let k = if !self.pre_reshape_qk_norm {